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 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_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_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_enqueue_two(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)
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, "text1")
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_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)
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)
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)
class TestScriptRunner(ScriptRunner):
"""Subclasses ScriptRunner to provide some testing features."""
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 _run_script_thread(self) -> None:
try:
super()._run_script_thread()
except BaseException as e:
self.script_thread_exceptions.append(e)
def _run_script(self, rerun_data: RerunData) -> None:
self.forward_msg_queue.clear()
super()._run_script(rerun_data)
def join(self) -> None:
"""Join the script_thread if it's running."""
if self._script_thread is not None:
self._script_thread.join()
def clear_forward_msgs(self) -> None:
"""Clear all messages from our ForwardMsgQueue."""
self.forward_msg_queue.clear()
def forward_msgs(self) -> List[ForwardMsg]:
"""Return all messages in our ForwardMsgQueue."""
return self.forward_msg_queue._queue
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: str, label: str) -> Optional[str]:
"""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
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()
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 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 _process_request_queue(self):
try:
super(TestScriptRunner, self)._process_request_queue()
except BaseException as e:
self.script_thread_exceptions.append(e)
def _run_script(self, rerun_data):
self.forward_msg_queue.clear()
super(TestScriptRunner, self)._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