Streamlit state management for React developers

This post assumes some basic familiarity with React and Streamlit. I do not profess expertise over either of them, but recently I have been struggling with a large Streamlit app and wanted to share a pattern which has been serving me well for keeping things tidy.

One of the nice things about React is the declarative mental model i.e. that UI is a function of state. Typically, a React functional component will look like this:

const MyComponent = (props: { initialState: string[] }) => {
  const { initialState } = props;

  // state management
  const [state, setState] = useState<string[]>(initialState);

  // usually an expensive computation
  const derivedState = useMemo(() => {
    return state.map(item => item.toUpperCase());
  }, [state]);

  // callback to handle events
  const handleClick = () => {
    setState(state.filter(item => item.length > 3));
  };

  return (
    <div>
      <button onClick={handleClick}>Click me</button>
      {derivedState}
    </div>
  );
};

At the top of the component, we have state related stuff (useState, useReducer, useContext, etc). Then, we’ll have logic to handle events and user interactions such as onClick callbacks, effects and so on. Finally at the bottom, we will return UI elements that are a function of the state.

With Streamlit, there is no opinionated orientation towards this pattern, but we can try to apply the same principles to make our code more readable/maintainable. The most important thing to know is that the entire Streamlit script behaves like a single React functional component. Every time you click a button or change some Streamlit session state, the entire script/app is re-run from top to bottom.

We can try to build out a useReducer style store and dispatch system to centralize our state, and bring in Pydantic for type-safe schemas:

from typing import Annotated, Literal, Union
from pydantic import BaseModel, Field

class Store(BaseModel):
    count: int = Field(default=0)
    first_name: str = Field(default="")

class IncrementCountAction(BaseModel):
    type: Literal["increment"] = "increment"

class SetFirstNameAction(BaseModel):
    type: Literal["set_first_name"] = "set_first_name"
    first_name: str

Action = Annotated[Union[IncrementCountAction, SetFirstNameAction], Field(discriminator="type")]

def reducer(state: Store, action: Action) -> Store:
  new_state = Store(**state.model_dump())
  if action.type == "increment":
    new_state.count += 1
  elif action.type == "set_first_name":
    new_state.first_name = action.first_name
  return new_state

def dispatch(action: Action):
  st.session_state.store = reducer(st.session_state.store, action)

def get_store() -> Store:
  return st.session_state.store

You could shove code like this into a store.py file and import it into your main Streamlit script. Basically, this code is setting up a single source of truth for the most important bits of your state. The store is just another field being set on Streamlit’s native session state which persists across page interactions.

Now, we can use this store in our main script:

from store import Store, dispatch, get_store
import streamlit as st

# state

if "store" not in st.session_state:
  st.session_state.store = Store()

# derived state (expensive computations)
@st.cache_data
def expensive_computation(count: int):
  return count * 2 # something much harder (like a network call)

# callbacks

def on_click_button():
  dispatch(IncrementCountAction())

def on_change_text_input():
  dispatch(SetFirstNameAction(first_name=st.session_state.bind_first_name))

# UI

st.button("Increment", on_click=on_click_button)

st.text_input("First name", key="bind_first_name", on_change=on_change_text_input, value=get_store().first_name)

if get_store().count > 5:
  st.write("Count is greater than 5")

if get_store().first_name:
  st.write(f"Hello, {get_store().first_name}")

st.write(f"Expensive computation: {expensive_computation(get_store().count)}")

Notice how the structure is similar to that of a React functional component.

Also, controlled inputs behave a little interestingly in Streamlit: each input gets some session state key assigned to it (which we can override to any key of our choice e.g. bind_first_name as I have done), and this can be used to access the value of the input via st.session_state.bind_first_name which can then be synchronized/dispatched to the central store. Meanwhile the initial value of the input can be retrieved from the store using get_store().

In summary, the approach described above is one where we do NOT use the return values of various inputs in an imperative manner (like e.g. first_name = st.text_input(...)). Instead we use callbacks to set some central state, and then use that to render UI accordingly.

Hope that’s helpful! If you’re interested in learning more, https://react.dev/learn/reacting-to-input-with-state is a well-written article from the React docs.