I created way simpler / cleaner session_state code with auto refresh

I created a very simple wrapper for session_state which allows you to do this:

import streamlit as st
from extensions import session_state_auto as ssa

st.title('Counter Example')
if not ssa.count:
    ssa.count = 0

st.write('Count = ', ssa.count)

ssa.count = st.number_input("Count", value=ssa.count)

if st.button('Increment'):
    ssa.count += 1

st.write('Count = ', ssa.count)

I’m new to streamlit & python, so i’m not sure why it was made so complicated with keys, and checking if keys exist, adding update functions, and on_clicks… why?

Anyways, this is ‘extensions.py’ I wrote:

import streamlit as st
from streamlit import session_state 

class SessionStateAutoClass:
    def __setattr__(self, name, value):
        if getattr(self, name) != value:
            session_state[name] = value
            st.experimental_rerun()

    def __getattr__(self, name):
        if name not in session_state:
            return None
        return session_state[name]

session_state_auto = SessionStateAutoClass()

Now tell me what I’m missing here and why this is bad.

3 Likes

Welcome to the forums!

Here’s what comes to mind in response to your inquiry (just to say what considerations come into play; not to say it’s bad or wrong):

One of my most common use cases for session state is doing things within callback functions. Streamlit does not like experimental reruns within a callback function since a rerun is already scheduled for the end of the callback. Hence, for my most frequent use case, I’d have to use the unmodified session state so as not to accidentally stuff in a rerun within the callback.

As for checking if keys are in session state, I am usually doing that at the top of the page to know whether or not I have to run some initializing script, so that ends up being a desirable feature for me. Streamlit is built around the idea that it reruns with every change, achieving efficiency with caching and rendering by incremental changes where it can. So the logic flow can feel a bit less linear, I think.

In your example, suppose I have a value in session state that was False, but initialized to True.

if not ssa.my_bool:
    ssa.my_bool = True

I’d have a problem with the above since the stored False value would always be overwritten. I could lengthen it to:

if ssa.my_bool is None:
    ssa.my_bool = True

We save maybe a few characters with standard aliasing in place, but what to do if I have a session state variable that might have value None with some other initializing value? In the end, the safest thing for me is to check for existence directly.

That being said, I do often set for brevity:

import streamlit as st
import st.session_state as ss

That’s just me though. :slight_smile: There are so many different use cases and logical flows that could exist, so if it’s helpful in your use case and those limitations don’t come up, that’s totally cool. :+1:

4 Likes

Thanks! That helps a lot. I had this sense of i’m forgetting something because I’m new. I’m used to more strongly typed languages, so the whole truthy thing is new for me… I wanted to embrace it!

I also had this version, which was more explicit in its first initialization, but then I simplified it too much maybe. Not sure if working with None is still fine if you don’t use it for truthy checks…(?)

I don’t do callbacks in my code…so is my way really bad? Why would I want callbacks?

import streamlit as st
from extensions import session_state_auto as ssa

st.title('Counter Example')
ssa.count = ssa.init(0)

st.write('Count = ', ssa.count)

ssa.count = st.number_input("Count", value=ssa.count)

if st.button('Increment'):
    ssa.count += 1

st.write('Count = ', ssa.count)
import streamlit as st
from streamlit import session_state 

class SessionStateAutoClass:
    def __setattr__(self, name, value):
        if type(value) is SessionValueWrapper:
            if name in session_state:
                return
            value = value.wrapped
        if self.__getattr__(name) != value:
            session_state[name] = value
            st.experimental_rerun()

    def __getattr__(self, name):
        if name not in session_state:
            return None          
        return session_state[name]
    
    def init(self, value):      
        return SessionValueWrapper(value)

class SessionValueWrapper:
    def __init__(self, wrapped):
        self.wrapped = wrapped

session_state_auto = SessionStateAutoClass()
1 Like

keys

session_state is a dictionary storing your values, hence you need keys to access the items in a dictionary.

Two things to note about this:

  1. This is not normal for a dictionary (because session_state is actually a dataclass wrapped around a dictionary), but session_state already provides access via the . notation, so you can access the key st.session_state["key"] by st.session_state.key.
  2. The attributes for your SessionStateAutoClass instance are also stored a dictionary - SessionStateAutoClass().__dict__.

So as far as the keys go, your class wrapper does not actually change anything. You still need the attribute name (key) to access the value, only instead of st.session_state.count you write ssa.count (if you want less typing, you can just alias session_state on import).

checking if keys exist

A key needs to exist in the dictionary if you want to access it. You could:

  1. Always assign the key to a default value first thing in your script
  2. Check if a key exists before reading/writing to it
  3. Wrap your access/assignment/update to the dictionary in try...except blocks

#1 would be nonsensical in the context of session_state. Why? Because session_state exists for you to preserve your variables between re-runs. If you always clear it, it will work like any other variable in the script (remember, streamlit re-runs the script on every interaction and normal variables are not stored, only session_state).

#2 and #3 are mechanistically different, but functionally equivalent ways of doing the same thing - in #2 you “look before you leap”, in #3 you “ask for forgiveness”.

Note that this does not apply just to dictionary keys, but to variables, attributes et cetera. Consider:

>>> class Foo():
...     def __init__(self):
...             self.bar = 42
...
>>> a = Foo()
>>> a.bar
42
>>> a.baz
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: 'Foo' object has no attribute 'baz'. Did you mean: 'bar'?

Anything in Python will throw an error if you reference it before assignment. You could, of course, construct a class which returns None if the attribute does not exist–as you did. This might make sense in some applications, but in Python, “explicit is better than implicit”. Having your application throw an exception if you try to access an object which does not exist is preferable to it silently returning None. In your implementation, you get to ignore the exception but then you still need to keep checking whether something is None or not, and if you forget, bad things might happen, and they might happen silently (for example if the value is used in a condition, as you already pointed out).

adding update functions, and on_clicks

First of all, on_click can call any function, not just those updating session_state. So the primary purpose of on_click and on_change callbacks is just to have something–anything–happen when you change a widget value.

Second, a callback can be used to immediately change the associated state without letting the script run until the end. So yes, changing a session_state value and calling an experimental_rerun() after a button press could be considered functionally equivalent to using a callback with a function changing the same value. But if you do this, and you want to abstract this code, you might want to put it inside in a function, and then instead of writing

if st.button("A button"):
    my_function()
    st.experimental_rerun()

You might want to write st.button("Button", on_click=my_function). This looks better, is easier to understand and pre-dates (I think) experimental_rerun, which is still experimental and could be removed.

The part which you find confusing (I think?) is where some people update session_state from a widget, in a callback from that widget. This is caused by a particular quirk of Streamlit, which really keeps track of two different session states:

  • one is the user session_state where you explicitly save values
  • the second allows you to save widget state to session_state by declaring the widget with a key parameter, it’s just syntactic sugar

The two are indistinguishable for most purposes, unless your widget vanishes at any point. The values you explicitly save to session_state are always preserved until they are explicitly removed or updated. If you declare a widget with a key parameter, the widget’s value is saved to session_state, but it’s removed if the widget is not rendered for some reason (pages, conditional widgets, etc.). Hence, some people will ignore the key parameter and have callbacks which explicitly save widget state to session_state, so that it persists. If your widgets always get rendered, you don’t have to update in callbacks, just put in the key parameter and the widget value will be saved in session_state automatically, like:

st.number_input("Count", key="count")

tell me what I’m missing here and why this is bad.

Mostly because if you wrap an API you do not fully understand, in a language you do not know very well, inside a wrapper class which changes its behavior, you might have a bad time when things go wrong. Streamlit API can be confusing at times, the last thing you want is to add another layer of confusion on top of it while you are figuring out how it works.

3 Likes