Retaining changes for Streamlit data editor when hiding or switching between widgets/pages

I created this simple python example below. I used streamlit together with pandas. This example has an editable dataframe in each selectbox "A" and "B". When I hit the selectbox "A", and edit the table: for example, add a new row as "a4" and "4" as value, then hit the selectbox "B" and come back to selectbox "A", the df1 goes back to original dataframe because the whole funct1 is rerun from the start. How can the edited dataframe information be stored so the edited dataframe infromation wouldn't be lost? I don't want to @st.cache_data it as I want the dataframe to be editable continuously.

import streamlit as st
import pandas as pd

page = st.sidebar.selectbox("Select: ", ("A","B"))

### Added code - but this doesn't work:
st.session_state['page'] = page
selected_app_mode = st.session_state.get('page')
app_mode_ix = 0
if selected_app_mode: 
    app_mode_ix = page.index(selected_app_mode)
page = st.sidebar.selectbox(page, index=app_mode_ix)
st.session_state['page'] = page
### End of added code

def funct1():
    df1 = pd.DataFrame({"col1": ["a1", "a2", "a3"], 
                       "Values": [1, 2, 3]})
    edited_df1 = st.experimental_data_editor(df1, num_rows="dynamic")
    return df1

def funct2():
    df2 = pd.DataFrame({"col1": ["b1", "b2", "b3"], 
                       "Values": [1, 2, 3]})
    edited_df1 = st.experimental_data_editor(df2, num_rows="dynamic")
    return df2

if  page == "A":
    funct1()
elif page == "B":
    funct2()

What I got (if I remove the added code):

df1
a1  1
a2  2
a3  3

Expected to get:

df1
a1  1
a2  2
a3  3
a4  4

Solution

Comments in the code below.

Something to keep in mind

The data editor is a little different than other widgets; you can't "store" its state directly. However, widgets lose their information when they disappear from the screen. This creates a problem.

For other widgets, you can save their value in session state (assigned to a different key than the widget's key) to keep their information while they are not displayed. When the widget comes back, you can assign it its previous state directly. However, because the data editor is the way it is, you can't directly save and assign its state. The best you can do is save the result of edits and then initialize a new editor that starts out where the previous one left off.

A caution

You don't want to feed a dataframe's edited result back into itself in real time. This will not work:

st.session_state.df = st.experimental_data_editor(st.session_state.df)

Such a pattern will cause the data editor to need each change entered twice to be reflected in the result. If an argument is changed in the creation of a widget, Streamlit thinks its a brand new widget and throws away any retained "memory" it had.

The solution

For each "page" you need to have two dataframes saved in session state: an original and an edited version. While on a page, you have a data editor based on the original and it saves the edited result directly into session state with each edit the user makes. When the page is changed, the edited version in session state is copied and overwrites the original one. Thus, when you return to the page, the data editor will start off where the last edit ended.

import streamlit as st
import pandas as pd

# Initialize session state with dataframes
# Include initialization of "edited" slots by copying originals
if 'df1' not in st.session_state:
    st.session_state.df1 = pd.DataFrame({
        "col1": ["a1", "a2", "a3"],
        "Values": [1, 2, 3]
    })
    st.session_state.edited_df1 = st.session_state.df1.copy()
    st.session_state.df2 = pd.DataFrame({
        "col1": ["b1", "b2", "b3"], 
        "Values": [1, 2, 3]
    })
    st.session_state.edited_df2 = st.session_state.df2.copy()

# Save edits by copying edited dataframes to "original" slots in session state
def save_edits():
    st.session_state.df1 = st.session_state.edited_df1.copy()
    st.session_state.df2 = st.session_state.edited_df2.copy()

# Sidebar to select page and commit changes upon selection
page = st.sidebar.selectbox("Select: ", ("A","B"), on_change=save_edits)

# Convenient shorthand notation
df1 = st.session_state.df1
df2 = st.session_state.df2

# Page functions commit edits in real time to "editied" slots in session state
def funct1():
    st.session_state.edited_df1 = st.experimental_data_editor(df1, num_rows="dynamic")
    return

def funct2():
    st.session_state.edited_df2 = st.experimental_data_editor(df2, num_rows="dynamic")
    return

if  page == "A":
    st.header("Page A")
    funct1()
elif page == "B":
    st.header("Page B")
    funct2()

PS. Strictly speaking, you can get away without the .copy() methods since the data editor is not performing any modification in place to the dataframe it's given. I just left them in as a kind of conceptual nod.

Edit: Further detailed explanation of the code per comment below

There are two pieces to focus on in the script:

page = st.sidebar.selectbox("Select: ", ("A","B"), on_change=save_edits)

and for each dataframe:

st.session_state.edited_df1 = st.experimental_data_editor(df1, num_rows="dynamic")

Say you have a page displaying the data for df1 for the user. If the user is editing the dataframe, then withe each edit:

User makes an edit
The value of the widget in session state is updated (we didn't use a manually assigned key, so you can't see this)
The page reloads
When the script gets to the widget again, it outputs the new state
This new output is saved to the edited_df1 key in session state.
Repeat 1-5 for however many edits the user does.
User changes to df2
on_change=save_edits executes before the new page load, hence st.session_state.edited_df1 is copied to st.session_state.df1 (same for df2 but it's trivial since they are the same)
Page reloads
User sees df2
Let's say the user immediately switches back to df1
Now the user sees the edited df1 because st.session_state.df1 was overwritten with the edited version when the user left that page/view