whichframe / SessionState.py
David Chuan-En Lin
init 7415fc4
1
"""Hack to add per-session state to Streamlit.
2
3
Works for Streamlit >= v0.65
4
5
Usage
6
-----
7
8
>>> import SessionState
9
>>>
10
>>> session_state = SessionState.get(user_name='', favorite_color='black')
11
>>> session_state.user_name
12
''
13
>>> session_state.user_name = 'Mary'
14
>>> session_state.favorite_color
15
'black'
16
17
Since you set user_name above, next time your script runs this will be the
18
result:
19
>>> session_state = get(user_name='', favorite_color='black')
20
>>> session_state.user_name
21
'Mary'
22
23
"""
24
25
import streamlit.report_thread as ReportThread
26
from streamlit.server.server import Server
27
28
29
class SessionState():
30
    """SessionState: Add per-session state to Streamlit."""
31
    def __init__(self, **kwargs):
32
        """A new SessionState object.
33
34
        Parameters
35
        ----------
36
        **kwargs : any
37
            Default values for the session state.
38
39
        Example
40
        -------
41
        >>> session_state = SessionState(user_name='', favorite_color='black')
42
        >>> session_state.user_name = 'Mary'
43
        ''
44
        >>> session_state.favorite_color
45
        'black'
46
47
        """
48
        for key, val in kwargs.items():
49
            setattr(self, key, val)
50
51
52
def get(**kwargs):
53
    """Gets a SessionState object for the current session.
54
55
    Creates a new object if necessary.
56
57
    Parameters
58
    ----------
59
    **kwargs : any
60
        Default values you want to add to the session state, if we're creating a
61
        new one.
62
63
    Example
64
    -------
65
    >>> session_state = get(user_name='', favorite_color='black')
66
    >>> session_state.user_name
67
    ''
68
    >>> session_state.user_name = 'Mary'
69
    >>> session_state.favorite_color
70
    'black'
71
72
    Since you set user_name above, next time your script runs this will be the
73
    result:
74
    >>> session_state = get(user_name='', favorite_color='black')
75
    >>> session_state.user_name
76
    'Mary'
77
78
    """
79
    # Hack to get the session object from Streamlit.
80
81
    session_id = ReportThread.get_report_ctx().session_id
82
    session_info = Server.get_current()._get_session_info(session_id)
83
84
    if session_info is None:
85
        raise RuntimeError('Could not get Streamlit session object.')
86
87
    this_session = session_info.session
88
89
    # Got the session object! Now let's attach some state into it.
90
91
    if not hasattr(this_session, '_custom_session_state'):
92
        this_session._custom_session_state = SessionState(**kwargs)
93
94
    return this_session._custom_session_state
95