plume-asr/src/plume/utils/st_rerun.py

46 lines
1.4 KiB
Python

try:
# Before Streamlit 0.65
from streamlit.ReportThread import get_report_ctx
from streamlit.server.Server import Server
from streamlit.ScriptRequestQueue import RerunData
from streamlit.ScriptRunner import RerunException
except ModuleNotFoundError:
# After Streamlit 0.65
from streamlit.report_thread import get_report_ctx
from streamlit.server.server import Server
from streamlit.script_request_queue import RerunData
from streamlit.script_runner import RerunException
def rerun():
"""Rerun a Streamlit app from the top!"""
widget_states = _get_widget_states()
raise RerunException(RerunData(widget_states))
def _get_widget_states():
# Hack to get the session object from Streamlit.
ctx = get_report_ctx()
session = None
current_server = Server.get_current()
if hasattr(current_server, '_session_infos'):
# Streamlit < 0.56
session_infos = Server.get_current()._session_infos.values()
else:
session_infos = Server.get_current()._session_info_by_id.values()
for session_info in session_infos:
if session_info.session.enqueue == ctx.enqueue:
session = session_info.session
if session is None:
raise RuntimeError(
"Oh noes. Couldn't get your Streamlit Session object"
"Are you doing something fancy with threads?"
)
# Got the session object!
return session._widget_states