Class property update does not get persisted when class instantiation is cached

I found a solution that leverages st.session_state.

What I do is the following:

  1. if A_instance is not in the session state yet (i.e., on the first run through the script), I instantiate A_instance in line 21. As such, I instantiate it only once.
  2. at the bottom of the script, I assign the st.session_state.A_instance with the instance of the class in line 33.
  3. When the script is executed again, A_instance is in st.session_state and the instantiation is skipped. Instead, I assign A_instance with st.session_state.A_instance in line 23, which has my instantiated class with the value of B retained.

Here is the code:

import streamlit as st
from datetime import datetime


class A:
    def __init__(self, instantiation_number):
        self.B = instantiation_number

    def update_B(self, number):
        self.B += number


@st.cache_data
def instantiate_a(instantiation_number):
    st.write("instantiating A at " + str(datetime.now()))

    return A(instantiation_number)


if "A_instance" not in st.session_state:
    A_instance = instantiate_a(0)
else:
    A_instance = st.session_state.A_instance

st.write(A_instance.B)

update_number = st.selectbox("Select a number", options=[1, 2, 3])

A_instance.update_B(update_number)

st.write(A_instance.B)

st.session_state.A_instance = A_instance