Python: Session State API で Streamlit をステートフルにする

これまで Streamlit で書いた Web アプリケーションは、基本的にステートレスだった。つまり、何らかのイベントが生じてアプリケーションのコードが再評価されると、ウィジェットを除くほとんどすべてのオブジェクトの状態はリセットされていた。アプリケーションをステートフルにする非公式なスニペットは一部で知られていたが、数行で使い始められるような気軽さはなかった。

そうした中、先日リリースされた Streamlit のバージョン 0.85 には、Session State API という機能が追加された。この API は、読んで字のごとく Streamlit の Web アプリケーションに限定的ながらステートを持たせることができる機能となっている。

docs.streamlit.io

今回は、追加された Session State API を触ってみることにした。

使った環境は次のとおり。

$ sw_vers                 
ProductName:    macOS
ProductVersion: 11.5.1
BuildVersion:   20G80
$ python -V                   
Python 3.9.6
$ pip list | grep -i streamlit 
streamlit                0.85.0

下準備

まずは肝心の Streamlit と、それ以外に可視化で使うデータセットを読み込むために Seaborn をインストールしておく。

$ pip install streamlit seaborn

ボタンを押すとカウンタが増減するサンプルコード

早速だけど、以下にカウンタの値をボタンに連動して増減させるサンプルコードを示す。 Session State API では、session_state という名前の辞書ライクなオブジェクトを扱う。このオブジェクトに格納したオブジェクト (以下、便宜的にセッション変数と呼ぶ) は、アプリケーションが再評価されても消えずに引き継がれる。セッション変数の値はウィジェットに追加されたコールバック関数の機能を介して更新する。以下では st.button() の on_change オプションにセッション変数の値を増減させるコールバック関数を登録している。

# -*- coding: utf-8 -*-

import streamlit as st


def main():
    # セッション変数が存在しないときは初期化する
    # ここでは 'counter' というセッション変数を作っている
    if 'counter' not in st.session_state:
        st.session_state['counter'] = 0

    # セッション変数の状態を表示する
    msg = f"Counter value: {st.session_state['counter']}"
    st.write(msg)

    # ボタンが押されたときに発火するコールバック
    def plus_one_clicks():
        # ボタンが押されたらセッション変数の値を増やす
        st.session_state['counter'] += 1
    # ボタンを作成するときにコールバックを登録しておく
    st.button(label='+1',
              on_click=plus_one_clicks)

    # ボタンが押されたらセッション変数の値を減らすバージョン
    def minus_one_clicks():
        st.session_state['counter'] -= 1
    st.button(label='-1',
              on_click=minus_one_clicks)

    # セッション変数の値をリセットするボタン
    def reset_clicks():
        st.session_state['counter'] = 0
    st.button(label='Reset',
              on_click=reset_clicks)


if __name__ == '__main__':
    main()

上記を保存したら Streamlit 経由で実行しよう。

$ streamlit run example.py

デフォルトでは自動で Web ブラウザが開くはず。開かない場合には以下でアクセスする。

$ open http://localhost:8501

すると、次のような WebUI が表示される。ボタンを押すと、それに連動してカウンタの値が増えたり減ったりする。

f:id:momijiame:20210728222510p:plain

これまで、ボタンをクリックするとイベントが生じてアプリケーションが再評価され、オブジェクトは一通りリセットされていた。しかし、Session State API を使うことで、それが回避できている。

Session State API を使う上での注意点は次のようなものがありそう。

(当たり前だけど) 存在しない変数 (辞書のキー) を使おうとすると例外になる
ブラウザをリロードすると変数はリセットされる
ページを複数のタブで開いたとしても変数は共有されない

データフレームのページネーションを実現するサンプルコード

続いては、もうちょっと実用的な例としてページネーションを実現してみる。以下のサンプルコードでは、タイタニックデータセットを読み込んで、それを 10 件ずつ表示するものになっている。表示している場所をセッション変数で管理することでページネーションが実現できる。

# -*- coding: utf-8 -*-

import math

import seaborn as sns
import streamlit as st


@st.cache
def load_dataset():
    """Titanic データセットを読み込む関数"""
    return sns.load_dataset('titanic')


def main():
    # データセットを読み込んで必要なページ数を計算する
    df = load_dataset()
    rows_per_page = 10
    total_pages = math.ceil(len(df) / rows_per_page)

    if 'page' not in st.session_state:
        st.session_state['page'] = 1

    left_col, center_col, right_col = st.beta_columns(3)

    # ページ数の増減ボタン
    with left_col:
        def minus_one_page():
            st.session_state['page'] -= 1
        if st.session_state['page'] > 1:
            st.button(label='<< Prev',
                      on_click=minus_one_page)

    with right_col:
        def plus_one_page():
            st.session_state['page'] += 1
        if st.session_state['page'] < total_pages:
            st.button(label='Next >>',
                      on_click=plus_one_page)

    # 現在のページ番号
    with center_col:
        st.write(f"Page: {st.session_state['page']} / {total_pages}")

    # ページ番号に応じた範囲のデータフレームを表示する
    start_iloc = (st.session_state['page'] - 1) * rows_per_page
    end_iloc = start_iloc + rows_per_page + 1
    st.write(df.iloc[start_iloc:end_iloc])


if __name__ == '__main__':
    main()