First stab at getting hooks to work

Author: Golf Player

nbclient/client.py

@@ -24,7 +24,7 @@
     CellExecutionComplete,
     CellExecutionError
 )
-from .util import run_sync, ensure_async
+from .util import run_sync, ensure_async, run_hook
 from .output_widget import OutputWidget
 
 
@@ -223,6 +223,35 @@ class NotebookClient(LoggingConfigurable):
 
     kernel_manager_class = Type(config=True, help='The kernel manager class to use.')
 
+    on_kernel_create = Any(
+        default_value=None,
+        allow_none=True,
+        help="""A callable which executes when the kernel is created.""",
+    ).tag(config=True)
+
+    on_cell_start = Any(
+        default_value=None,
+        allow_none=True,
+        help="""A callable which executes before a cell is executed.""",
+    ).tag(config=True)
+
+    on_cell_complete = Any(
+        default_value=None,
+        allow_none=True,
+        help=dedent(
+            """
+            A callable which executes after a cell execution is complete. It is 
+            called even when a cell results in a failure.
+            """
+        ),
+    ).tag(config=True)
+
+    on_cell_error = Any(
+        default_value=None,
+        allow_none=True,
+        help="""A callable which executes when a cell execution results in an error.""",
+    ).tag(config=True)
+
     @default('kernel_manager_class')
     def _kernel_manager_class_default(self):
         """Use a dynamic default to avoid importing jupyter_client at startup"""
@@ -378,6 +407,7 @@ async def async_start_new_kernel_client(self, **kwargs):
 
         kernel_id = await ensure_async(self.km.start_kernel(extra_arguments=self.extra_arguments,
                                        **kwargs))
+        run_hook(self.on_kernel_create, kernel_id)
 
         # if self.km is not a KernelManager, it's probably a MultiKernelManager
         try:
@@ -661,14 +691,15 @@ def _passed_deadline(self, deadline):
             return True
         return False
 
-    def _check_raise_for_error(self, cell, exec_reply):
+    def _check_raise_for_error(self, cell, cell_index, exec_reply):
         cell_allows_errors = self.allow_errors or "raises-exception" in cell.metadata.get(
             "tags", []
         )
 
-        if self.force_raise_errors or not cell_allows_errors:
-            if (exec_reply is not None) and exec_reply['content']['status'] == 'error':
-                raise CellExecutionError.from_cell_and_msg(cell, exec_reply['content'])
+        if (exec_reply is not None) and exec_reply['content']['status'] == 'error':
+            run_hook(self.on_cell_error, cell, cell_index)
+            if self.force_raise_errors or not cell_allows_errors:
+                    raise CellExecutionError.from_cell_and_msg(cell, exec_reply['content'])
 
     async def async_execute_cell(self, cell, cell_index, execution_count=None, store_history=True):
         """
@@ -712,6 +743,7 @@ async def async_execute_cell(self, cell, cell_index, execution_count=None, store
             cell['metadata']['execution'] = {}
 
         self.log.debug("Executing cell:\n%s", cell.source)
+        run_hook(self.on_cell_start, cell, cell_index)
         parent_msg_id = await ensure_async(
             self.kc.execute(
                 cell.source,
@@ -744,7 +776,7 @@ async def async_execute_cell(self, cell, cell_index, execution_count=None, store
 
         if execution_count:
             cell['execution_count'] = execution_count
-        self._check_raise_for_error(cell, exec_reply)
+        self._check_raise_for_error(cell, cell_index, exec_reply)
         self.nb['cells'][cell_index] = cell
         return cell
 

nbclient/util.py

@@ -89,3 +89,10 @@ async def ensure_async(obj):
         return result
     # obj doesn't need to be awaited
     return obj
+
+
+def run_hook(hook, *args):
+    if hook is None:
+        return
+    loop = asyncio.get_event_loop()
+    loop.call_soon(hook, *args)