streamlitで変数の状態を維持する

※注意
本記事はstreamlit version 0.84.0以前の記事となります。
Streamlit version 0.84.0にてSession State機能がリリースされたため、本記事で行っているようなSession管理を自前実装する必要はなくなりました。
詳細は公式ドキュメントを参照ください: https://docs.streamlit.io/en/stable/add_state_app.html

こんにちは。
streamlitを使っているとき、その挙動の特殊さに苦労することが多いです。
今回は変数の状態維持についてどのように実現するかを書きます。

github: https://github.com/irisu-inwl/streamlit-state-test

現象

ボタンを押すとcountが増加、減退するアプリケーションを作ろうとします。
以下のようにコードを書くと、一回目の施行は成功しますが、以降ボタンを押しても変数は+1, -1から動きません。

import streamlit as st


def main():
    st.title('Counter App')
    count = 0
    increment_count = st.button('count +')
    decrement_count = st.button('count -')
    if increment_count:
        count += 1
    if decrement_count:
        count -= 1

    st.write(f'count: {count}')


if __name__ == '__main__':
    main()

run app

docker run -itd -v $(pwd)/src:/opt/streamlit/src -p 8080:8501 --name streamlit-state streamlit-state streamlit run src/wrong_app.py

結果:

streamlitの実行順序

streamlitの実行順序をおさらいします。
streamlitの公式ドキュメントでは、streamlitの挙動として、アプリのrendering後にユーザーによるイベント発火でスクリプトを再実行するとあります。
そのため、ボタンを押した際に毎回count = 0が実行され変数が初期化されてしまい、+1,-1の状態となってしまいます。

引用元:

対処

では、変数の状態を維持するためにはどうするか、というと、以下のstreamlit communityの記事にてセッションの状態を管理する方法が記されております。

引用されてるコードにある_SessionStateクラスを利用して、先ほどのアプリを以下のように改良します。

import streamlit as st
from src.session import _get_state


state = _get_state()

if state.count == None:
    state.count = 0


def main():
    st.title('Counter App')

    increment_count = st.button('count +')
    decrement_count = st.button('count -')
    if increment_count:
        state.count += 1
    if decrement_count:
        state.count -= 1

    st.write(f'count: {state.count}')


if __name__ == '__main__':
    main()

まず、_SessionStateクラス, _get_state(), _get_session()をsession.pyに書きます。(こんな感じ)
そして、アプリケーション実行時にcountを初期化します。その際に_SessionStateにcount変数の状態を保存するためstate.countを参照することで、アプリケーションのイベントが発生で初期化されないようにします。

以降アプリケーションコードで状態を管理したい変数はstateから参照すれば以下のように状態が保持されます。

おわり

streamlitで変数の状態を管理する方法を紹介しました。
streamlitはめちゃめちゃ便利ですが、使う上で、その挙動について知ることと、何が出来て何ができないかを知ることが大切だと思います。（任意のことにも言えますが）
では。