feat: implement Output widget that mimics a frontend

This is a port of voila-dashboards/voila#91 and subsequent fixes.

nbclient/client.py

@@ -1,5 +1,6 @@
-import datetime
 import base64
+import collections
+import datetime
 from textwrap import dedent
 
 from async_generator import asynccontextmanager
@@ -22,6 +23,7 @@
     CellExecutionError
 )
 from .util import run_sync, ensure_async
+from .output_widget import OutputWidget
 
 
 def timestamp():
@@ -307,6 +309,11 @@ def reset_execution_trackers(self):
         self._display_id_map = {}
         self.widget_state = {}
         self.widget_buffers = {}
+        # maps to list of hooks, where the last is used, this is used
+        # to support nested use of output widgets.
+        self.output_hook_stack = collections.defaultdict(list)
+        # our front-end mimicing Output widgets
+        self.output_widget_objects = {}
 
     def start_kernel_manager(self):
         """Creates a new kernel manager.
@@ -787,6 +794,14 @@ def process_message(self, msg, cell, cell_index):
     def output(self, outs, msg, display_id, cell_index):
         msg_type = msg['msg_type']
 
+        parent_msg_id = msg['parent_header'].get('msg_id')
+        if self.output_hook_stack[parent_msg_id]:
+            # if we have a hook registered, it will overrride our
+            # default output behaviour (e.g. OutputWidget)
+            hook = self.output_hook_stack[parent_msg_id][-1]
+            hook.output(outs, msg, display_id, cell_index)
+            return
+
         try:
             out = output_from_msg(msg)
         except ValueError:
@@ -812,6 +827,15 @@ def output(self, outs, msg, display_id, cell_index):
 
     def clear_output(self, outs, msg, cell_index):
         content = msg['content']
+
+        parent_msg_id = msg['parent_header'].get('msg_id')
+        if self.output_hook_stack[parent_msg_id]:
+            # if we have a hook registered, it will overrride our
+            # default clear_output behaviour (e.g. OutputWidget)
+            hook = self.output_hook_stack[parent_msg_id][-1]
+            hook.clear_output(outs, msg, cell_index)
+            return
+
         if content.get('wait'):
             self.log.debug('Wait to clear output')
             self.clear_before_next_output = True
@@ -832,6 +856,23 @@ def handle_comm_msg(self, outs, msg, cell_index):
             self.widget_state.setdefault(content['comm_id'], {}).update(data['state'])
             if 'buffer_paths' in data and data['buffer_paths']:
                 self.widget_buffers[content['comm_id']] = self._get_buffer_data(msg)
+        # There are cases where we need to mimic a frontend, to get similar behaviour as
+        # when using the Output widget from Jupyter lab/notebook
+        if msg['msg_type'] == 'comm_open' and msg['content'].get('target_name') == 'jupyter.widget':
+            content = msg['content']
+            data = content['data']
+            state = data['state']
+            comm_id = msg['content']['comm_id']
+            if state['_model_module'] == '@jupyter-widgets/output' and state['_model_name'] == 'OutputModel':
+                self.output_widget_objects[comm_id] = OutputWidget(comm_id, state, self.kc, self)
+        elif msg['msg_type'] == 'comm_msg':
+            content = msg['content']
+            data = content['data']
+            if 'state' in data:
+                state = data['state']
+                comm_id = msg['content']['comm_id']
+                if comm_id in self.output_widget_objects:
+                    self.output_widget_objects[comm_id].set_state(state)
 
     def _serialize_widget_state(self, state):
         """Serialize a widget state, following format in @jupyter-widgets/schema."""
@@ -856,6 +897,22 @@ def _get_buffer_data(self, msg):
             )
         return encoded_buffers
 
+    def register_output_hook(self, msg_id, hook):
+        """Registers an override object that handles output/clear_output instead.
+
+        Multiple hooks can be registered, where the last one will be used (stack based)
+        """
+        # mimics
+        # https://jupyterlab.github.io/jupyterlab/services/interfaces/kernel.ikernelconnection.html#registermessagehook
+        self.output_hook_stack[msg_id].append(hook)
+
+    def remove_output_hook(self, msg_id, hook):
+        """Unregisters an override object that handles output/clear_output instead"""
+        # mimics
+        # https://jupyterlab.github.io/jupyterlab/services/interfaces/kernel.ikernelconnection.html#removemessagehook
+        removed_hook = self.output_hook_stack[msg_id].pop()
+        assert removed_hook == hook
+
 
 def execute(nb, cwd=None, km=None, **kwargs):
     """Execute a notebook's code, updating outputs within the notebook object.

nbclient/output_widget.py

@@ -0,0 +1,76 @@
+from ipykernel.jsonutil import json_clean
+from nbformat.v4 import output_from_msg
+
+
+class OutputWidget:
+    """This class mimics a front end output widget"""
+    def __init__(self, comm_id, state, kernel_client, executor):
+        self.comm_id = comm_id
+        self.state = state
+        self.kernel_client = kernel_client
+        self.executor = executor
+        self.topic = ('comm-%s' % self.comm_id).encode('ascii')
+        self.outputs = self.state['outputs']
+        self.clear_before_next_output = False
+
+    def clear_output(self, outs, msg, cell_index):
+        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):
+        state = {'outputs': self.outputs}
+        msg = {'method': 'update', 'state': state, 'buffer_paths': []}
+        self.send(msg)
+
+    def _publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
+        """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=None, metadata=None, buffers=None):
+        self._publish_msg('comm_msg', data=data, metadata=metadata, buffers=buffers)
+
+    def output(self, outs, msg, display_id, cell_index):
+        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):
+        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