Expander does not keep its state if something is displayed above it

I am having some troubles with st.expander. In the example below, if the second expander is opened and then the button is clicked, on the next run, the first expander is opened while the second one is closed. Any idea what is causing that?

import streamlit as st

# Load the data
button = st.button('Click me')

# Display message
container = st.container()
if button:
    with container:
        st.markdown('Data loaded')

# Expanders
with st.expander('Expander 1'):
    st.markdown('Expander 1X')
with st.expander('Expander 2'):
    st.markdown('Exander 2X')
with st.expander('Expander 3'):
    st.markdown('Exander 3X')

I’ve recorded a quick video to explain the problem

Emmanuel

It seems that setting container as st.empty first and then calling container() solve my problem. Hopefully that helps people with the same problem.

import streamlit as st

# Load the data
button = st.button('Click me')

# Display message
container = st.empty()
if button:
    with container.container():
        st.write('Data loaded')

# Expanders
with st.expander('Expander 1'):
    st.markdown('Expander 1X')
with st.expander('Expander 2'):
    st.markdown('Exander 2X')
with st.expander('Expander 3'):
    st.markdown('Exander 3X')