Galuh commited on
Commit
8ae1004
1 Parent(s): 0e919bc

Create SessionState.py

Browse files
Files changed (1) hide show
  1. SessionState.py +107 -0
SessionState.py ADDED
@@ -0,0 +1,107 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Hack to add per-session state to Streamlit.
2
+ Usage
3
+ -----
4
+ >>> import SessionState
5
+ >>>
6
+ >>> session_state = SessionState.get(user_name='', favorite_color='black')
7
+ >>> session_state.user_name
8
+ ''
9
+ >>> session_state.user_name = 'Mary'
10
+ >>> session_state.favorite_color
11
+ 'black'
12
+ Since you set user_name above, next time your script runs this will be the
13
+ result:
14
+ >>> session_state = get(user_name='', favorite_color='black')
15
+ >>> session_state.user_name
16
+ 'Mary'
17
+ """
18
+ try:
19
+ import streamlit.ReportThread as ReportThread
20
+ from streamlit.server.Server import Server
21
+ except Exception:
22
+ # Streamlit >= 0.65.0
23
+ import streamlit.report_thread as ReportThread
24
+ from streamlit.server.server import Server
25
+
26
+
27
+ class SessionState(object):
28
+ def __init__(self, **kwargs):
29
+ """A new SessionState object.
30
+ Parameters
31
+ ----------
32
+ **kwargs : any
33
+ Default values for the session state.
34
+ Example
35
+ -------
36
+ >>> session_state = SessionState(user_name='', favorite_color='black')
37
+ >>> session_state.user_name = 'Mary'
38
+ ''
39
+ >>> session_state.favorite_color
40
+ 'black'
41
+ """
42
+ for key, val in kwargs.items():
43
+ setattr(self, key, val)
44
+
45
+
46
+ def get(**kwargs):
47
+ """Gets a SessionState object for the current session.
48
+ Creates a new object if necessary.
49
+ Parameters
50
+ ----------
51
+ **kwargs : any
52
+ Default values you want to add to the session state, if we're creating a
53
+ new one.
54
+ Example
55
+ -------
56
+ >>> session_state = get(user_name='', favorite_color='black')
57
+ >>> session_state.user_name
58
+ ''
59
+ >>> session_state.user_name = 'Mary'
60
+ >>> session_state.favorite_color
61
+ 'black'
62
+ Since you set user_name above, next time your script runs this will be the
63
+ result:
64
+ >>> session_state = get(user_name='', favorite_color='black')
65
+ >>> session_state.user_name
66
+ 'Mary'
67
+ """
68
+ # Hack to get the session object from Streamlit.
69
+
70
+ ctx = ReportThread.get_report_ctx()
71
+
72
+ this_session = None
73
+
74
+ current_server = Server.get_current()
75
+ if hasattr(current_server, '_session_infos'):
76
+ # Streamlit < 0.56
77
+ session_infos = Server.get_current()._session_infos.values()
78
+ else:
79
+ session_infos = Server.get_current()._session_info_by_id.values()
80
+
81
+ for session_info in session_infos:
82
+ s = session_info.session
83
+ if (
84
+ # Streamlit < 0.54.0
85
+ (hasattr(s, '_main_dg') and s._main_dg == ctx.main_dg)
86
+ or
87
+ # Streamlit >= 0.54.0
88
+ (not hasattr(s, '_main_dg') and s.enqueue == ctx.enqueue)
89
+ or
90
+ # Streamlit >= 0.65.2
91
+ (not hasattr(s, '_main_dg') and s._uploaded_file_mgr == ctx.uploaded_file_mgr)
92
+ ):
93
+ this_session = s
94
+
95
+ if this_session is None:
96
+ raise RuntimeError(
97
+ "Oh noes. Couldn't get your Streamlit Session object. "
98
+ 'Are you doing something fancy with threads?')
99
+
100
+ # Got the session object! Now let's attach some state into it.
101
+
102
+ if not hasattr(this_session, '_custom_session_state'):
103
+ this_session._custom_session_state = SessionState(**kwargs)
104
+
105
+ return this_session._custom_session_state
106
+
107
+ __all__ = ['get']