Following good feedback on my post in Preserving state across sidebar pages I decided to implement a prototype implementation.
This is basically the same as the proposed SessionState but takes advantage of python typing to let the user define a class/dataclass to define state and therefore take advantage of type checking and code completion in modern editors (tested in Pycharm).
It’s basically a session_state.py file with the implementation logic (adapted from https://gist.github.com/tvst/036da038ab3e999a64497f42de966a92)
from typing import Callable, TypeVar
from streamlit import ReportThread
from streamlit.server.Server import Server
T = TypeVar('T')
# noinspection PyProtectedMember
def get_state(setup_func: Callable[..., T], **kwargs) -> T:
ctx = ReportThread.get_report_ctx()
session = None
session_infos = Server.get_current()._session_infos.values()
for session_info in session_infos:
if session_info.session._main_dg == ctx.main_dg:
session = session_info.session
if session is None:
raise RuntimeError(
"Oh noes. Couldn't get your Streamlit Session object"
'Are you doing something fancy with threads?')
# Got the session object! Now let's attach some state into it.
if not getattr(session, '_custom_session_state', None):
session._custom_session_state = setup_func(**kwargs)
return session._custom_session_state
And here is a test program using it to increment a counter when a button is pressed.
import streamlit as st
from session_state import get_state
class MyState:
a: int
b: str
def __init__(self, a: int, b: str):
self.a = a
self.b = b
def setup(a: int, b: str) -> MyState:
print('Running setup')
return MyState(a, b)
state = get_state(setup, a=3, b='bbb')
st.title('Test state')
if st.button('Increment'):
state.a = state.a + 1
print(f'Incremented to {state.a}')
st.text(state.a)
In summary, get_state takes a setup function that will only be executed exactly once (this is an improvement over SessionState as far as I can tell because whatever is in the get was always executed). The setup_function needs to return an object, and whatever object type it returns will also be the return type of get_state. If it’s the first time we are calling get_state then we call the setup function that initializes our state (optionally with keyword arguments). The following times we call it it returns the instance of the state.
Any feedback is appreciated.
