class SessionData:
"""
Contains parameters related to running a script, and also houses
the ForwardMsgQueue that is used to deliver messages to a connected browser.
"""
script_path: str
script_folder: str
name: str
command_line: str
script_run_id: str
_browser_queue: ForwardMsgQueue
def __init__(self, script_path: str, command_line: str):
"""Constructor.
Parameters
----------
script_path : str
Path of the Python file from which this app is generated.
command_line : string
Command line as input by the user
"""
basename = os.path.basename(script_path)
self.script_path = os.path.abspath(script_path)
self.script_folder = os.path.dirname(self.script_path)
self.name = str(os.path.splitext(basename)[0])
# The browser queue contains messages that haven't yet been
# delivered to the browser. Periodically, the server flushes
# this queue and delivers its contents to the browser.
self._browser_queue = ForwardMsgQueue()
self.script_run_id = generate_new_id()
self.command_line = command_line
def enqueue(self, msg):
self._browser_queue.enqueue(msg)
def clear(self):
self._browser_queue.clear()
def flush_browser_queue(self):
"""Clears our browser queue and returns the messages it contained.
The Server calls this periodically to deliver new messages
to the browser associated with this session.
Returns
-------
list[ForwardMsg]
The messages that were removed from the queue and should
be delivered to the browser.
"""
return self._browser_queue.flush()
def __init__(self, script_name):
"""Initializes the ScriptRunner for the given script_name"""
# DeltaGenerator deltas will be enqueued into self.forward_msg_queue.
self.forward_msg_queue = ForwardMsgQueue()
def enqueue_fn(msg):
self.forward_msg_queue.enqueue(msg)
self.maybe_handle_execution_control_request()
self.script_request_queue = ScriptRequestQueue()
script_path = os.path.join(os.path.dirname(__file__), "test_data",
script_name)
super(TestScriptRunner, self).__init__(
session_id="test session id",
session_data=SessionData(script_path, "test command line"),
enqueue_forward_msg=enqueue_fn,
client_state=ClientState(),
session_state=SessionState(),
request_queue=self.script_request_queue,
uploaded_file_mgr=UploadedFileManager(),
)
# Accumulates uncaught exceptions thrown by our run thread.
self.script_thread_exceptions = []
# Accumulates all ScriptRunnerEvents emitted by us.
self.events = []
self.event_data = []
def record_event(event, **kwargs):
self.events.append(event)
self.event_data.append(kwargs)
self.on_event.connect(record_event, weak=False)
def __init__(self, script_path: str, command_line: str):
"""Constructor.
Parameters
----------
script_path : str
Path of the Python file from which this app is generated.
command_line : string
Command line as input by the user
"""
basename = os.path.basename(script_path)
self.script_path = os.path.abspath(script_path)
self.script_folder = os.path.dirname(self.script_path)
self.name = str(os.path.splitext(basename)[0])
# The browser queue contains messages that haven't yet been
# delivered to the browser. Periodically, the server flushes
# this queue and delivers its contents to the browser.
self._browser_queue = ForwardMsgQueue()
self.script_run_id = generate_new_id()
self.command_line = command_line
class DeltaGeneratorTestCase(unittest.TestCase):
def setUp(self, override_root=True):
self.forward_msg_queue = ForwardMsgQueue()
self.override_root = override_root
self.orig_report_ctx = None
self.new_script_run_ctx = ScriptRunContext(
session_id="test session id",
enqueue=self.forward_msg_queue.enqueue,
query_string="",
session_state=SessionState(),
uploaded_file_mgr=UploadedFileManager(),
)
if self.override_root:
self.orig_report_ctx = get_script_run_ctx()
add_script_run_ctx(threading.current_thread(),
self.new_script_run_ctx)
self.app_session = FakeAppSession()
def tearDown(self):
self.clear_queue()
if self.override_root:
add_script_run_ctx(threading.current_thread(),
self.orig_report_ctx)
def get_message_from_queue(self, index=-1):
"""Get a ForwardMsg proto from the queue, by index.
Returns
-------
ForwardMsg
"""
return self.forward_msg_queue._queue[index]
def get_delta_from_queue(self, index=-1):
"""Get a Delta proto from the queue, by index.
Returns
-------
Delta
"""
deltas = self.get_all_deltas_from_queue()
return deltas[index]
def get_all_deltas_from_queue(self):
"""Return all the delta messages in our ForwardMsgQueue"""
return [
msg.delta for msg in self.forward_msg_queue._queue
if msg.HasField("delta")
]
def clear_queue(self):
self.forward_msg_queue._clear()
def setUp(self, override_root=True):
self.forward_msg_queue = ForwardMsgQueue()
self.override_root = override_root
self.orig_report_ctx = None
self.new_script_run_ctx = ScriptRunContext(
session_id="test session id",
enqueue=self.forward_msg_queue.enqueue,
query_string="",
session_state=SessionState(),
uploaded_file_mgr=UploadedFileManager(),
)
if self.override_root:
self.orig_report_ctx = get_script_run_ctx()
add_script_run_ctx(threading.current_thread(),
self.new_script_run_ctx)
self.app_session = FakeAppSession()
def test_simple_enqueue(self):
rq = ForwardMsgQueue()
self.assertTrue(rq.is_empty())
rq.enqueue(NEW_SESSION_MSG)
self.assertFalse(rq.is_empty())
queue = rq.flush()
self.assertTrue(rq.is_empty())
self.assertEqual(len(queue), 1)
self.assertTrue(queue[0].new_session.config.allow_run_on_save)
def test_replace_element(self):
rq = ForwardMsgQueue()
self.assertTrue(rq.is_empty())
rq.enqueue(NEW_SESSION_MSG)
TEXT_DELTA_MSG1.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 0)
rq.enqueue(TEXT_DELTA_MSG1)
TEXT_DELTA_MSG2.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 0)
rq.enqueue(TEXT_DELTA_MSG2)
queue = rq.flush()
self.assertEqual(len(queue), 2)
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 0),
queue[1].metadata.delta_path)
self.assertEqual(queue[1].delta.new_element.text.body, "text2")
def __init__(self, script_name: str):
"""Initializes the ScriptRunner for the given script_name"""
# DeltaGenerator deltas will be enqueued into self.forward_msg_queue.
self.forward_msg_queue = ForwardMsgQueue()
main_script_path = os.path.join(os.path.dirname(__file__), "test_data",
script_name)
super().__init__(
session_id="test session id",
session_data=SessionData(main_script_path, "test command line"),
client_state=ClientState(),
session_state=SessionState(),
uploaded_file_mgr=UploadedFileManager(),
initial_rerun_data=RerunData(),
)
# Accumulates uncaught exceptions thrown by our run thread.
self.script_thread_exceptions: List[BaseException] = []
# Accumulates all ScriptRunnerEvents emitted by us.
self.events: List[ScriptRunnerEvent] = []
self.event_data: List[Any] = []
def record_event(sender: Optional[ScriptRunner],
event: ScriptRunnerEvent, **kwargs) -> None:
# Assert that we're not getting unexpected `sender` params
# from ScriptRunner.on_event
assert (sender is None
or sender == self), "Unexpected ScriptRunnerEvent sender!"
self.events.append(event)
self.event_data.append(kwargs)
# Send ENQUEUE_FORWARD_MSGs to our queue
if event == ScriptRunnerEvent.ENQUEUE_FORWARD_MSG:
forward_msg = kwargs["forward_msg"]
self.forward_msg_queue.enqueue(forward_msg)
self.on_event.connect(record_event, weak=False)
def test_replace_element(self):
"""Enqueuing an element with the same delta_path as another element
already in the queue should replace the original element.
"""
rq = ForwardMsgQueue()
self.assertTrue(rq.is_empty())
rq.enqueue(NEW_SESSION_MSG)
TEXT_DELTA_MSG1.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 0)
rq.enqueue(TEXT_DELTA_MSG1)
TEXT_DELTA_MSG2.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 0)
rq.enqueue(TEXT_DELTA_MSG2)
queue = rq.flush()
self.assertEqual(2, len(queue))
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 0),
queue[1].metadata.delta_path)
self.assertEqual("text2", queue[1].delta.new_element.text.body)
def test_enqueue_three(self):
"""Enqueue 3 ForwardMsgs."""
rq = ForwardMsgQueue()
self.assertTrue(rq.is_empty())
rq.enqueue(NEW_SESSION_MSG)
TEXT_DELTA_MSG1.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 0)
rq.enqueue(TEXT_DELTA_MSG1)
TEXT_DELTA_MSG2.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 1)
rq.enqueue(TEXT_DELTA_MSG2)
queue = rq.flush()
self.assertEqual(3, len(queue))
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 0),
queue[1].metadata.delta_path)
self.assertEqual("text1", queue[1].delta.new_element.text.body)
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 1),
queue[2].metadata.delta_path)
self.assertEqual("text2", queue[2].delta.new_element.text.body)
def test_multiple_containers(self):
"""Deltas should only be coalesced if they're in the same container"""
rq = ForwardMsgQueue()
self.assertTrue(rq.is_empty())
rq.enqueue(NEW_SESSION_MSG)
def enqueue_deltas(container: int, path: Tuple[int, ...]):
# We deep-copy the protos because we mutate each one
# multiple times.
msg = copy.deepcopy(TEXT_DELTA_MSG1)
msg.metadata.delta_path[:] = make_delta_path(container, path, 0)
rq.enqueue(msg)
msg = copy.deepcopy(DF_DELTA_MSG)
msg.metadata.delta_path[:] = make_delta_path(container, path, 1)
rq.enqueue(msg)
msg = copy.deepcopy(ADD_ROWS_MSG)
msg.metadata.delta_path[:] = make_delta_path(container, path, 1)
rq.enqueue(msg)
enqueue_deltas(RootContainer.MAIN, ())
enqueue_deltas(RootContainer.SIDEBAR, (0, 0, 1))
def assert_deltas(container: int, path: Tuple[int, ...], idx: int):
# Text delta
self.assertEqual(make_delta_path(container, path, 0),
queue[idx].metadata.delta_path)
self.assertEqual("text1", queue[idx].delta.new_element.text.body)
# Dataframe delta
self.assertEqual(make_delta_path(container, path, 1),
queue[idx + 1].metadata.delta_path)
col0 = queue[
idx + 1].delta.new_element.data_frame.data.cols[0].int64s.data
col1 = queue[
idx + 1].delta.new_element.data_frame.data.cols[1].int64s.data
self.assertEqual([0, 1, 2], col0)
self.assertEqual([10, 11, 12], col1)
# add_rows delta
self.assertEqual(make_delta_path(container, path, 1),
queue[idx + 2].metadata.delta_path)
ar_col0 = queue[idx +
2].delta.add_rows.data.data.cols[0].int64s.data
ar_col1 = queue[idx +
2].delta.add_rows.data.data.cols[1].int64s.data
self.assertEqual([3, 4, 5], ar_col0)
self.assertEqual([13, 14, 15], ar_col1)
queue = rq.flush()
self.assertEqual(7, len(queue))
assert_deltas(RootContainer.MAIN, (), 1)
assert_deltas(RootContainer.SIDEBAR, (0, 0, 1), 4)
def test_dont_replace_block(self, other_msg: ForwardMsg):
"""add_block deltas should never be replaced/composed because they can
have dependent deltas later in the queue."""
rq = ForwardMsgQueue()
self.assertTrue(rq.is_empty())
ADD_BLOCK_MSG.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 0)
other_msg.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 0)
# Delta messages should not replace `add_block` deltas with the
# same delta_path.
rq.enqueue(ADD_BLOCK_MSG)
rq.enqueue(other_msg)
queue = rq.flush()
self.assertEqual(len(queue), 2)
self.assertEqual(queue[0], ADD_BLOCK_MSG)
self.assertEqual(queue[1], other_msg)
def test_cached_st_function_warning(self, _, cache_decorator, call_stack):
"""Ensure we properly warn when st.foo functions are called
inside a cached function.
"""
forward_msg_queue = ForwardMsgQueue()
orig_report_ctx = get_script_run_ctx()
add_script_run_ctx(
threading.current_thread(),
ScriptRunContext(
session_id="test session id",
enqueue=forward_msg_queue.enqueue,
query_string="",
session_state=SessionState(),
uploaded_file_mgr=None,
),
)
with patch.object(call_stack,
"_show_cached_st_function_warning") as warning:
st.text("foo")
warning.assert_not_called()
@cache_decorator
def cached_func():
st.text("Inside cached func")
cached_func()
warning.assert_called_once()
warning.reset_mock()
# Make sure everything got reset properly
st.text("foo")
warning.assert_not_called()
# Test warning suppression
@cache_decorator(suppress_st_warning=True)
def suppressed_cached_func():
st.text("No warnings here!")
suppressed_cached_func()
warning.assert_not_called()
# Test nested st.cache functions
@cache_decorator
def outer():
@cache_decorator
def inner():
st.text("Inside nested cached func")
return inner()
outer()
warning.assert_called_once()
warning.reset_mock()
# Test st.cache functions that raise errors
with self.assertRaises(RuntimeError):
@cache_decorator
def cached_raise_error():
st.text("About to throw")
raise RuntimeError("avast!")
cached_raise_error()
warning.assert_called_once()
warning.reset_mock()
# Make sure everything got reset properly
st.text("foo")
warning.assert_not_called()
# Test st.cache functions with widgets
@cache_decorator
def cached_widget():
st.button("Press me!")
cached_widget()
warning.assert_called_once()
warning.reset_mock()
# Make sure everything got reset properly
st.text("foo")
warning.assert_not_called()
add_script_run_ctx(threading.current_thread(), orig_report_ctx)
def test_add_rows_rerun(self):
rq = ForwardMsgQueue()
self.assertTrue(rq.is_empty())
rq.enqueue(NEW_SESSION_MSG)
# Simulate rerun
for i in range(2):
TEXT_DELTA_MSG1.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 0)
rq.enqueue(TEXT_DELTA_MSG1)
DF_DELTA_MSG.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 1)
rq.enqueue(DF_DELTA_MSG)
ADD_ROWS_MSG.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 1)
rq.enqueue(ADD_ROWS_MSG)
queue = rq.flush()
self.assertEqual(5, len(queue))
# Text delta
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 0),
queue[1].metadata.delta_path)
self.assertEqual("text1", queue[1].delta.new_element.text.body)
# Dataframe delta
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 1),
queue[2].metadata.delta_path)
col0 = queue[2].delta.new_element.data_frame.data.cols[0].int64s.data
col1 = queue[2].delta.new_element.data_frame.data.cols[1].int64s.data
self.assertEqual([0, 1, 2], col0)
self.assertEqual([10, 11, 12], col1)
# First add_rows delta
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 1),
queue[3].metadata.delta_path)
ar_col0 = queue[3].delta.add_rows.data.data.cols[0].int64s.data
ar_col1 = queue[3].delta.add_rows.data.data.cols[1].int64s.data
self.assertEqual([3, 4, 5], ar_col0)
self.assertEqual([13, 14, 15], ar_col1)
# Second add_rows delta
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 1),
queue[4].metadata.delta_path)
ar_col0 = queue[4].delta.add_rows.data.data.cols[0].int64s.data
ar_col1 = queue[4].delta.add_rows.data.data.cols[1].int64s.data
self.assertEqual([3, 4, 5], ar_col0)
self.assertEqual([13, 14, 15], ar_col1)
class TestScriptRunner(ScriptRunner):
"""Subclasses ScriptRunner to provide some testing features."""
def __init__(self, script_name):
"""Initializes the ScriptRunner for the given script_name"""
# DeltaGenerator deltas will be enqueued into self.forward_msg_queue.
self.forward_msg_queue = ForwardMsgQueue()
self.script_request_queue = ScriptRequestQueue()
main_script_path = os.path.join(os.path.dirname(__file__), "test_data",
script_name)
super(TestScriptRunner, self).__init__(
session_id="test session id",
session_data=SessionData(main_script_path, "test command line"),
enqueue_forward_msg=self.forward_msg_queue.enqueue,
client_state=ClientState(),
session_state=SessionState(),
request_queue=self.script_request_queue,
uploaded_file_mgr=UploadedFileManager(),
)
# Accumulates uncaught exceptions thrown by our run thread.
self.script_thread_exceptions = []
# Accumulates all ScriptRunnerEvents emitted by us.
self.events: List[ScriptRunnerEvent] = []
self.event_data: List[Any] = []
def record_event(sender: Optional[ScriptRunner],
event: ScriptRunnerEvent, **kwargs):
# Assert that we're not getting unexpected `sender` params
# from ScriptRunner.on_event
assert (sender is None
or sender == self), "Unexpected ScriptRunnerEvent sender!"
self.events.append(event)
self.event_data.append(kwargs)
self.on_event.connect(record_event, weak=False)
def enqueue_rerun(self, argv=None, widget_states=None, query_string=""):
self.script_request_queue.enqueue(
ScriptRequest.RERUN,
RerunData(widget_states=widget_states, query_string=query_string),
)
def enqueue_stop(self):
self.script_request_queue.enqueue(ScriptRequest.STOP)
def enqueue_shutdown(self):
self.script_request_queue.enqueue(ScriptRequest.SHUTDOWN)
def _run_script_thread(self):
try:
super()._run_script_thread()
except BaseException as e:
self.script_thread_exceptions.append(e)
def _run_script(self, rerun_data):
self.forward_msg_queue.clear()
super()._run_script(rerun_data)
def join(self):
"""Joins the run thread, if it was started"""
if self._script_thread is not None:
self._script_thread.join()
def clear_deltas(self):
"""Clear all delta messages from our ForwardMsgQueue"""
self.forward_msg_queue.clear()
def deltas(self) -> List[Delta]:
"""Return the delta messages in our ForwardMsgQueue"""
return [
msg.delta for msg in self.forward_msg_queue._queue
if msg.HasField("delta")
]
def elements(self) -> List[Element]:
"""Return the delta.new_element messages in our ForwardMsgQueue."""
return [delta.new_element for delta in self.deltas()]
def text_deltas(self) -> List[str]:
"""Return the string contents of text deltas in our ForwardMsgQueue"""
return [
element.text.body for element in self.elements()
if element.WhichOneof("type") == "text"
]
def get_widget_id(self, widget_type, label):
"""Returns the id of the widget with the specified type and label"""
for delta in self.deltas():
new_element = getattr(delta, "new_element", None)
widget = getattr(new_element, widget_type, None)
widget_label = getattr(widget, "label", None)
if widget_label == label:
return widget.id
return None
def test_add_rows_rerun(self):
rq = ForwardMsgQueue()
self.assertTrue(rq.is_empty())
rq.enqueue(NEW_SESSION_MSG)
# Simulate rerun
for i in range(2):
TEXT_DELTA_MSG1.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 0)
rq.enqueue(TEXT_DELTA_MSG1)
DF_DELTA_MSG.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 1)
rq.enqueue(DF_DELTA_MSG)
ADD_ROWS_MSG.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 1)
rq.enqueue(ADD_ROWS_MSG)
queue = rq.flush()
self.assertEqual(len(queue), 3)
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 0),
queue[1].metadata.delta_path)
self.assertEqual(queue[1].delta.new_element.text.body, "text1")
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 1),
queue[2].metadata.delta_path)
col0 = queue[2].delta.new_element.data_frame.data.cols[0].int64s.data
col1 = queue[2].delta.new_element.data_frame.data.cols[1].int64s.data
self.assertEqual(col0, [0, 1, 2, 3, 4, 5])
self.assertEqual(col1, [10, 11, 12, 13, 14, 15])
def test_simple_add_rows(self):
"""'add_rows' messages should behave as expected."""
rq = ForwardMsgQueue()
self.assertTrue(rq.is_empty())
rq.enqueue(NEW_SESSION_MSG)
TEXT_DELTA_MSG1.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 0)
rq.enqueue(TEXT_DELTA_MSG1)
DF_DELTA_MSG.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 1)
rq.enqueue(DF_DELTA_MSG)
ADD_ROWS_MSG.metadata.delta_path[:] = make_delta_path(
RootContainer.MAIN, (), 1)
rq.enqueue(ADD_ROWS_MSG)
queue = rq.flush()
self.assertEqual(4, len(queue))
# Text delta
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 0),
queue[1].metadata.delta_path)
self.assertEqual("text1", queue[1].delta.new_element.text.body)
# Dataframe delta
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 1),
queue[2].metadata.delta_path)
df_col0 = queue[2].delta.new_element.data_frame.data.cols[
0].int64s.data
df_col1 = queue[2].delta.new_element.data_frame.data.cols[
1].int64s.data
self.assertEqual([0, 1, 2], df_col0)
self.assertEqual([10, 11, 12], df_col1)
# AddRows delta
self.assertEqual(make_delta_path(RootContainer.MAIN, (), 1),
queue[3].metadata.delta_path)
ar_col0 = queue[3].delta.add_rows.data.data.cols[0].int64s.data
ar_col1 = queue[3].delta.add_rows.data.data.cols[1].int64s.data
self.assertEqual([3, 4, 5], ar_col0)
self.assertEqual([13, 14, 15], ar_col1)