mirror of
https://github.com/aykhans/AzSuicideDataVisualization.git
synced 2025-04-21 10:15:45 +00:00
112 lines
3.9 KiB
Python
112 lines
3.9 KiB
Python
from typing import Any, Dict, List, Optional
|
|
|
|
from jupyter_client.client import KernelClient
|
|
from nbformat.v4 import output_from_msg
|
|
|
|
from .jsonutil import json_clean
|
|
|
|
|
|
class OutputWidget:
|
|
"""This class mimics a front end output widget"""
|
|
|
|
def __init__(
|
|
self, comm_id: str, state: Dict[str, Any], kernel_client: KernelClient, executor: Any
|
|
) -> None:
|
|
|
|
self.comm_id: str = comm_id
|
|
self.state: Dict[str, Any] = state
|
|
self.kernel_client: KernelClient = kernel_client
|
|
self.executor = executor
|
|
self.topic: bytes = ('comm-%s' % self.comm_id).encode('ascii')
|
|
self.outputs: List = self.state['outputs']
|
|
self.clear_before_next_output: bool = False
|
|
|
|
def clear_output(self, outs: List, msg: Dict, cell_index: int) -> None:
|
|
|
|
self.parent_header = msg['parent_header']
|
|
content = msg['content']
|
|
if content.get('wait'):
|
|
self.clear_before_next_output = True
|
|
else:
|
|
self.outputs = []
|
|
# sync back the state to the kernel
|
|
self.sync_state()
|
|
if hasattr(self.executor, 'widget_state'):
|
|
# sync the state to the nbconvert state as well, since that is used for testing
|
|
self.executor.widget_state[self.comm_id]['outputs'] = self.outputs
|
|
|
|
def sync_state(self) -> None:
|
|
state = {'outputs': self.outputs}
|
|
msg = {'method': 'update', 'state': state, 'buffer_paths': []}
|
|
self.send(msg)
|
|
|
|
def _publish_msg(
|
|
self,
|
|
msg_type: str,
|
|
data: Optional[Dict] = None,
|
|
metadata: Optional[Dict] = None,
|
|
buffers: Optional[List] = None,
|
|
**keys: Any
|
|
) -> None:
|
|
"""Helper for sending a comm message on IOPub"""
|
|
data = {} if data is None else data
|
|
metadata = {} if metadata is None else metadata
|
|
content = json_clean(dict(data=data, comm_id=self.comm_id, **keys))
|
|
msg = self.kernel_client.session.msg(
|
|
msg_type, content=content, parent=self.parent_header, metadata=metadata
|
|
)
|
|
self.kernel_client.shell_channel.send(msg)
|
|
|
|
def send(
|
|
self,
|
|
data: Optional[Dict] = None,
|
|
metadata: Optional[Dict] = None,
|
|
buffers: Optional[List] = None,
|
|
) -> None:
|
|
|
|
self._publish_msg('comm_msg', data=data, metadata=metadata, buffers=buffers)
|
|
|
|
def output(self, outs: List, msg: Dict, display_id: str, cell_index: int) -> None:
|
|
|
|
if self.clear_before_next_output:
|
|
self.outputs = []
|
|
self.clear_before_next_output = False
|
|
self.parent_header = msg['parent_header']
|
|
output = output_from_msg(msg)
|
|
|
|
if self.outputs:
|
|
# try to coalesce/merge output text
|
|
last_output = self.outputs[-1]
|
|
if (
|
|
last_output['output_type'] == 'stream'
|
|
and output['output_type'] == 'stream'
|
|
and last_output['name'] == output['name']
|
|
):
|
|
last_output['text'] += output['text']
|
|
else:
|
|
self.outputs.append(output)
|
|
else:
|
|
self.outputs.append(output)
|
|
self.sync_state()
|
|
if hasattr(self.executor, 'widget_state'):
|
|
# sync the state to the nbconvert state as well, since that is used for testing
|
|
self.executor.widget_state[self.comm_id]['outputs'] = self.outputs
|
|
|
|
def set_state(self, state: Dict) -> None:
|
|
if 'msg_id' in state:
|
|
msg_id = state.get('msg_id')
|
|
if msg_id:
|
|
self.executor.register_output_hook(msg_id, self)
|
|
self.msg_id = msg_id
|
|
else:
|
|
self.executor.remove_output_hook(self.msg_id, self)
|
|
self.msg_id = msg_id
|
|
|
|
def handle_msg(self, msg: Dict) -> None:
|
|
content = msg['content']
|
|
comm_id = content['comm_id']
|
|
assert comm_id == self.comm_id
|
|
data = content['data']
|
|
if 'state' in data:
|
|
self.set_state(data['state'])
|