Session_state seems to be dropping some data

stratovate · November 10, 2022, 12:30pm

Summary

I’m trying out session_state in a simple test case. It’s largely working, but on one particular action it clears one of its attributes, and I can’t see why.

Steps to reproduce

Code snippet:

import streamlit as st

st.session_state

if 'running_query' not in st.session_state:
    st.session_state.running_query = False
if 'prefix' not in st.session_state:
    st.session_state.prefix = 'ABC'

st.session_state

def run_query():
    st.session_state.running_query = True

def change_prefix():
    st.session_state.running_query = False

if not st.session_state.running_query:
    st.write("Not running query")
    title = st.text_input('Code prefix:', key='prefix')
    st.button('Run query', on_click=run_query)
else:
    st.write('Running query with:', st.session_state.prefix)
    st.button('Change prefix', on_click=change_prefix)

If applicable, please provide the steps we should take to reproduce the error or specified behavior.

Expected behavior:

I should be able to set the prefix by typing “test” in the input box to replace the default “ABC”, hit enter to apply, then ‘run query’ and see a message saying “Running query with test”.

Then click “change prefix” and the default in the Code prefix input box should be “test”.

Actual behavior:

When clicking “change prefix” in the final step described above, instead the default goes back to “ABC”.

Debug info

Streamlit version: 1.14.0
Python version: 3.9.13
Using Conda
OS version: macOS 13.0.1
Browser version: Chrome Version 107.0.5304.110 (Official Build) (arm64)pytho

Requirements file

name: streamlit
channels:

defaults
conda-forge
dependencies:
pip
matplotlib=3.5.0
numpy=1.21.2
plotly=5.6.0
python=3.9.13
pip:
- pandas==1.5.0
- pip==22.2.2
- pyathena==2.5.1
- sqlalchemy==1.4.32
- streamlit==1.14.0
- toml==0.10.2

Additional information

I put in the two magic calls to display st.session_state and it seems clear that when the ‘change prefix’ button is pressed, the session_state data has been lost - running_query is set to False but prefix is nowhere to be seen, so the default init statement sets it back to ABC. Why is this happening? I want the input box to be pre-populated with whatever the user put in last time, so they can edit it.

To be clear, this is just a test rig for me to understand session_state, so I’m not looking for workarounds - I want to understand why it’s not holding on to the state as I think it should be?

mathcatsand · November 10, 2022, 3:14pm

Your code is updating the prefix fine for me. Are you hitting enter to submit the text before clicking the Run query button? If you type in text but click the button before submitting the text (with enter), it won’t register the text there.

You can use a form to get around the need to ‘submit the text’ before clicking the button the run the query.

import streamlit as st

st.session_state

if 'running_query' not in st.session_state:
    st.session_state.running_query = False
if 'prefix' not in st.session_state:
    st.session_state.prefix = 'ABC'

st.session_state

def run_query():
    st.session_state.running_query = True

def change_prefix():
    st.session_state.running_query = False

change_state = st.form(key='my_form')
with change_state:
    if not st.session_state.running_query:
        st.write("Not running query")
        title = st.text_input('Code prefix:', key='prefix')
        st.form_submit_button('Run query', on_click=run_query)
    else:
        st.write('Running query with:', st.session_state.prefix)
        st.form_submit_button('Change prefix', on_click=change_prefix)

stratovate · November 10, 2022, 3:46pm

Strange… yes, I am pressing enter to apply the input. It does indeed pick up my input in the next stage (showing “Running query with: test” or whatever I’ve entered) - this isn’t where I’m seeing a problem.

The problem comes when returning to edit the prefix by clicking on the ‘change prefix’ at this point, which then always gives me the ABC default in the input box, rather than ‘test’ or whatever I was most recently using.

To be clear, what I would expect to happen when pressing ‘change prefix’ is that session_state.prefix retains its value. But instead it gets dropped and reset to the default.

Can you confirm @mathcatsand that you were seeing the correct behaviour after clicking ‘change prefix’? In which case that’s weird, given the same setup and the same code…

Btw, I’ve tried changing to use a form as you suggest (v helpful hint - thanks) but it doesn’t make any difference to this session_state issue.

mathcatsand · November 10, 2022, 4:08pm

I see your question. My understanding is that since the key is associated with a widget (the gets destroyed), it is getting removed from session_state. You can buffer that by copying it over to a different key that Streamlit won’t delete (by way of association to a disappearing input), or utilize the fact the input box is getting destroyed and use value to set a default from something kept in session_state.

Here’s the first suggestion as an example, with an extra print of session_state so you can see what happens from the callback running to the next page load.

import streamlit as st

st.session_state

if 'running_query' not in st.session_state:
    st.session_state.running_query = False
if 'prefix' not in st.session_state:
    st.session_state.prefix = 'ABC'
if 'prefix_memory' not in st.session_state:
    st.session_state.prefix_memory = 'ABC'

st.session_state

def run_query():
    st.session_state
    st.write('^^ This is run_query running, just before page reload.')
    # This extracts the prefixs into a different slot in session state before
    # page reload so it can survive the destruction of the input widget which
    # will be gone on next page load.
    st.session_state.prefix_memory = st.session_state.prefix
    st.session_state.running_query = True

def change_prefix():
    st.session_state
    st.write('^^ This is change_prefix running, just before page reload.')
    st.session_state.running_query = False

change_state = st.form(key='my_form')
with change_state:
    if not st.session_state.running_query:
        st.write("Not running query")
        title = st.text_input('Code prefix:', key='prefix', value=st.session_state.prefix_memory)
        st.form_submit_button('Run query', on_click=run_query)
    else:
        st.write('Running query with:', st.session_state.prefix)
        st.form_submit_button('Change prefix', on_click=change_prefix)

stratovate · November 10, 2022, 4:32pm

Ah, thanks @mathcatsand … got it.

Seems strange behaviour though, that using the session_state property (prefix here) as the key to the widget is what the documentation seems to recommend, but I can’t see it anywhere that it makes mention of the fact that it gets dropped when the widget goes.

You suggest two methods - the first, as per your code, makes sense. But I’d like to understand the second option too, and can’t get it to work. Sorry - feel pretty stupid coming back again with such a basic problem, but I don’t seem to have got my head round how streamlit handles these inputs.

Here’s what I’ve done to try to implement your option 2:

import streamlit as st

st.session_state

if 'running_query' not in st.session_state:
    st.session_state.running_query = False
if 'prefix' not in st.session_state:
    st.session_state.prefix = 'ABC'

st.session_state

def run_query(prefix):
    st.session_state.running_query = True
    st.session_state.prefix = prefix

def change_prefix():
    st.session_state.running_query = False

change_state = st.form(key='my_form')
with change_state:
    if not st.session_state.running_query:
        st.write("Not running query")
        edited_prefix = st.text_input('Code prefix', value=st.session_state.prefix)
        st.form_submit_button('Run query', on_click=run_query, args=(edited_prefix,))
    else:
        st.write('Running query with:', st.session_state.prefix)
        st.form_submit_button('Change prefix', on_click=change_prefix)

So this gives me a different problem: the edited prefix is never picked up, so whatever I type in the input box, the parameter supplied to the run_query function is always ‘ABC’ and the prefix doesn’t change.

What am I missing?

mathcatsand · November 10, 2022, 4:42pm

Sorry, I think I mashed the two solutions a bit from playing around with the code. To clarify, here’s two (slightly) different solutions:

import streamlit as st

st.session_state

if 'running_query' not in st.session_state:
    st.session_state.running_query = False
if 'prefix_memory' not in st.session_state:
    st.session_state.prefix_memory = 'ABC'

st.session_state

def run_query():
    st.session_state.prefix_memory = st.session_state.prefix
    st.session_state.running_query = True

def change_prefix():
    st.session_state.running_query = False

change_state = st.form(key='my_form')
with change_state:
    if not st.session_state.running_query:
        st.write("Not running query")
        # We don't do anything manual to assign value to st.session_state.prefix
        # We are only using it as a key to access the information during the callback.
        # The prefix key goes away when the text input widget goes away, 
        # so this widget is recreated each time it shows up
        title = st.text_input('Code prefix:', key='prefix', value=st.session_state.prefix_memory)
        st.form_submit_button('Run query', on_click=run_query)
    else:
        st.write('Running query with:', st.session_state.prefix)
        st.form_submit_button('Change prefix', on_click=change_prefix)

import streamlit as st

st.session_state

if 'running_query' not in st.session_state:
    st.session_state.running_query = False
if 'prefix_memory' not in st.session_state:
    st.session_state.prefix_memory = 'ABC'
# Here we protect against the destruction by re-initializing from prefix_memory
if 'prefix' not in st.session_state:
    st.session_state.prefix = st.session_state.prefix_memory

st.session_state

def run_query():
    st.session_state.prefix_memory = st.session_state.prefix
    st.session_state.running_query = True

def change_prefix():
    st.session_state.running_query = False

change_state = st.form(key='my_form')
with change_state:
    if not st.session_state.running_query:
        st.write("Not running query")
        title = st.text_input('Code prefix:', key='prefix')
        st.form_submit_button('Run query', on_click=run_query)
    else:
        st.write('Running query with:', st.session_state.prefix)
        st.form_submit_button('Change prefix', on_click=change_prefix)

stratovate · November 10, 2022, 4:59pm

Thanks very much @mathcatsand

mathcatsand · November 10, 2022, 5:07pm

And just to clarify a bit on what you were working with, in this snippet, edited_prefix will get the value from what was in the input box on page load, not from what was just submitted. Hence, the callback function will never see the “new” entry. If you need to access a new submission within a callback, you have to do it from st.session_state.<widget key> since the order of operations is

[[0. If a widget is set to “output” to a variable, it gets the pre-existing/default value on page load]]

Store information using the widget’s key
Run the callback
Reload the page at which point the widget will “output” to any variable assignment based on the new page load (if the widget persisted, it will remember the last submission)