How to control work flow in Streamlit with session state

Streamlit is super easy to use to build some demo quickly. Streamlit is an open-source Python library that allows developers to easily create interactive web applications for data science and machine learning projects. With Streamlit, you can create a web-based dashboard or application that lets users interact with your data, visualizations, and machine learning models without the need for extensive coding or web development knowledge

But if you want to more complex work flow controls, you might expect some unexpected behaviors.
For example, a nested buttons, where users click the first button, do something, then continue to click the second button, and do other things.
Such as the following picture shows:

As shown in the above picture, the goal we want to achieve is: First, user will clicks button1, then a random number got displayed, then the button2 showsup, and users continue to click button2, and should
see a new random generated.

You might think the streamlit code is as easy as the following:

import streamlit as st
import random



button1 = st.button('Step1')
if button1 == True:
        
    show_number1 = random.random()
    st.write("random number generated from step 1: {}".format(show_number1))
    


    button2 = st.button('Step2')
    if button2 == True:
        show_number2 = random.random()
        st.write("random number generated from step 2: {}".format(show_number2))
    
    

For the above code, when you click button1, it works as you would expect. But when you click button2, the page restarts, and you only see button1 again.
This is not a bug, but what’s supposed to work with Streamlit, probably for a good reason.

So if we want to achieve our original goal, we need to leverage session state to store some several things. One thing to notice, once we click the second button, the app will restart anyway,
we can’t really change that, we just need to cache some knowledge of what already happened, and use those information to guide the work flow.

So information we need to cache are:
(1) if each of the buttons are already clicked
(2) the random number generated after the first button click, so it doesn’t get regnerated again during the app restart.

The new code is like the following:

import streamlit as st
import random


## init states
if 'button1_clicked' not in st.session_state:
    st.session_state.button1_clicked = False
    
if 'button2_clicked' not in st.session_state:
    st.session_state.button2_clicked = False
    
if 'show_number1' not in st.session_state:
    st.session_state.show_number1 = 0

    
def callback1():
    st.session_state.button1_clicked = True
    
def callback2():
    st.session_state.button2_clicked = True


button1 = st.button('Step1',on_click=callback1)
if button1 == True or st.session_state.button2_clicked == True:
    
    if button1 == False and st.session_state.button2_clicked  == True:
        show_number1 = st.session_state.show_number1
    else:
        show_number1 = random.random()
        
    st.write("random number generated from step 1: {}".format(show_number1))
    st.session_state.show_number1 = show_number1

    button2 = st.button('Step2',on_click=callback2)
    if button2 == True:
        show_number2 = random.random()
        st.write("random number generated from step 2: {}".format(show_number2))