ReturnOnCode Blog

Write readable Streamlit code using React's useState model

Engineering - Bas Dunn - August 2024

At ReturnOnCode, we use Streamlit to showcase our AI applications to clients. While we love Streamlit for its simplicity, we're also big fans of NextJS/React for building robust, production-ready front-end systems. Although these tools serve different purposes, they share quite a few similarities. Both are front-end frameworks, and their approaches to interacting with the framework have a lot in common.

While working on an AI application using Streamlit, we realized that we could borrow some features from React to enhance the readability of our Streamlit code.

Streamlining Session State Management in Streamlit Using React's useState Model

For more advanced Streamlit (and React) applications, managing session state is crucial. In Streamlit, you use the session_state feature, while in React, you have useState.

In Streamlit, managing session state typically looks like the following code block. We find this approach a bit cumbersome. There are too many lines required to initialize session_state, and defining a setter function for every variable seems redundant. Also, the repetition of the variable name 'message' at least three times in our code feels unnecessarily verbose.

import streamlit as st

# Initialize session_state and define setter functions
if 'message' not in st.session_state:
    st.session_state['message'] = '...'
    
def set_message(message: str):
    st.session_state['message'] = message

st.write(st.session_state['message'])
st.button("say hi!", on_click=set_message, args=("Hi!",))

A Better Approach: Instead of writing all these lines of code, we suggest borrowing the useState model from React. As shown in the code block below, we've abstracted away the initialization and setter function into a separate module. This keeps our Streamlit application code clean and concise.

import streamlit as st
from use_state import useState

# Initialize session_state and define setter functions
message, set_message = useState('message', initial_value="...")

st.write(message)
st.button("say hi!", on_click=set_message, args=("Hi!",))

Creating the useState Module for Streamlit

Here's the simple Python code to create your own useState module. Feel free to reach out if you have any suggestions for improvements!

# use_state.py
import streamlit as st

def useState(name: str, initial_value=None):
    """
    React-like session_state wrapper.
    
    Returns:
        Tuple containing:
        - The value
        - A function to set the value
    """
    
    if name not in st.session_state:
        st.session_state[name] = initial_value  
        
    def setter(v):
        st.session_state[name] = v
        
    return st.session_state[name], setter

Want to Learn How We Can Use These Tools to Boost Your Business?

Reach out to us using the 'Meet us' button.