Alternative implementation of session state

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.

Slightly better example using a dataclass and not passing kwargs:

import streamlit as st

from session_state import get_state

@dataclass
class MyState:
    a: int
    b: str

def setup() -> MyState:
    print('Running setup')
    return MyState(a=3, b='bbb')

state = get_state(setup)

st.title('Test state')
if st.button('Increment'):
    state.a = state.a + 1
    print(f'Incremented to {state.a}')
st.text(state.a)

I’ve been using this setup for the past few days with great success. It really makes buttons more useful when state can be stored.

2 Likes

Hi @biellls

I really like it.

Is it possible to do an additional simplification of the api?

I imagine you could actual “hide” the setup function in you session_state.py file and change the public api from state = get_state(setup) to state=get_state(a=3, b='bbb').

Marc

Hi @Marc, thanks for the feedback. What you’re proposing is basically what was implemented in the original gist if I understood correctly. Let me try to explain why I think my way is an improvement over that for two reasons:

  1. Arguments get evaluated every time there’s a re-run, so if you want to open a database connection and you do state=get_state(db=psycopg2.connect("dbname=test user=postgres password=secret")) you’ll be opening it lots of times.
    The reason I added a setup function is that this way we can guarantee that the code inside it runs exactly once, which is what a setup should do.
    Ideally we’d have that in a closure/multi-line lambda but python doesn’t allow those (and regular lambdas are untyped) so I had to make it a function.

  2. You lose typing information. The original get_state is an object that gets added instance variables dynamically, whereas in mine I created a class/dataclass. Since that class/dataclass is the return of my setup function, get_state can infer that it’s returning that class/dataclass and then type checking/code complete works in the editor.

Makes sense. Thanks.

Very cool. We’ve been thinking about this as well and will post some possible designs. :+1: