refactor: make comm_open handler extensible

Author: maartenbreddels

nbclient/client.py

@@ -306,6 +306,10 @@ def __init__(self, nb, km=None, **kw):
                 'OutputModel': OutputWidget
             }
         }
+        # comm_open_handlers should return an object with a .handle_msg(msg) method or None
+        self.comm_open_handlers = {
+            'jupyter.widget': self.on_comm_open_jupyter_widget
+        }
 
     def reset_execution_trackers(self):
         """Resets any per-execution trackers.
@@ -318,7 +322,7 @@ def reset_execution_trackers(self):
         # to support nested use of output widgets.
         self.output_hook_stack = collections.defaultdict(list)
         # our front-end mimicing Output widgets
-        self.widget_objects = {}
+        self.comm_objects = {}
 
     def start_kernel_manager(self):
         """Creates a new kernel manager.
@@ -863,23 +867,17 @@ def handle_comm_msg(self, outs, msg, cell_index):
                 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']
+        if msg['msg_type'] == 'comm_open':
+            handler = self.comm_open_handlers.get(msg['content'].get('target_name'))
             comm_id = msg['content']['comm_id']
-            module = self.widget_registry.get(state['_model_module'])
-            if module:
-               widget_class = module.get(state['_model_name'])
-               self.widget_objects[comm_id] = widget_class(comm_id, state, self.kc, self)
+            comm_object = handler(msg)
+            if comm_object:
+                self.comm_objects[comm_id] = comm_object
         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.widget_objects:
-                    self.widget_objects[comm_id].set_state(state)
+            comm_id = msg['content']['comm_id']
+            if comm_id in self.comm_objects:
+                self.comm_objects[comm_id].handle_msg(msg)
 
     def _serialize_widget_state(self, state):
         """Serialize a widget state, following format in @jupyter-widgets/schema."""
@@ -920,6 +918,17 @@ def remove_output_hook(self, msg_id, hook):
         removed_hook = self.output_hook_stack[msg_id].pop()
         assert removed_hook == hook
 
+    def on_comm_open_jupyter_widget(self, msg):
+        content = msg['content']
+        data = content['data']
+        state = data['state']
+        comm_id = msg['content']['comm_id']
+        module = self.widget_registry.get(state['_model_module'])
+        if module:
+            widget_class = module.get(state['_model_name'])
+            if widget_class:
+                return widget_class(comm_id, state, self.kc, self)
+
 
 def execute(nb, cwd=None, km=None, **kwargs):
     """Execute a notebook's code, updating outputs within the notebook object.

nbclient/output_widget.py

@@ -75,3 +75,11 @@ def set_state(self, state):
             else:
                 self.executor.remove_output_hook(self.msg_id, self)
                 self.msg_id = msg_id
+
+    def handle_msg(self, msg):
+        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'])