Create session_state.py

session_state.py

@@ -0,0 +1,28 @@
+class _SessionState:
+    def __init__(self, **kwargs):
+        for key, val in kwargs.items():
+            setattr(self, key, val)
+
+
+def get(**kwargs):
+    from streamlit.report_thread import get_report_ctx
+
+    session_id = get_report_ctx().session_id
+    session = _get_session(session_id)
+    session_state = session.setdefault("state", _SessionState(**kwargs))
+    return session_state
+
+
+def _get_session(session_id):
+    from streamlit.server.server import Server
+
+    server = Server.get_current()
+    if not hasattr(server, "_session_infos"):
+        raise RuntimeError("Oh noes. Couldn't get your Streamlit Session object.")
+
+    session_infos = server._session_infos.values()
+    for session_info in session_infos:
+        if session_info.session._main_dg.report_ctx.session_id == session_id:
+            return session_info.session
+
+    raise RuntimeError("Oh noes. Couldn't get your Streamlit Session object.")