Add basic hooks during execution

This will enable tracking of execution process without subclassing the
way papermill does.

Author: Golf Player

nbclient/client.py

@@ -26,7 +26,7 @@
     DeadKernelError,
 )
 from .output_widget import OutputWidget
-from .util import ensure_async, run_sync
+from .util import ensure_async, run_sync, run_hook
 
 
 def timestamp(msg: Optional[Dict] = None) -> str:
@@ -261,6 +261,45 @@ class NotebookClient(LoggingConfigurable):
 
     kernel_manager_class: KernelManager = Type(config=True, help='The kernel manager class to use.')
 
+    on_execution_start: t.Optional[t.Callable] = Any(
+        default_value=None,
+        allow_none=True,
+        help=dedent("""
+        Called after the kernel manager and kernel client are setup, and cells
+        are about to execute.
+        Called with kwargs `kernel_id`.
+        """),
+    ).tag(config=True)
+
+    on_cell_start: t.Optional[t.Callable] = Any(
+        default_value=None,
+        allow_none=True,
+        help=dedent("""
+        A callable which executes before a cell is executed.
+        Called with kwargs `cell`, and `cell_index`.
+        """),
+    ).tag(config=True)
+
+    on_cell_complete: t.Optional[t.Callable] = 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.
+        Called with kwargs `cell`, and `cell_index`.
+        """),
+    ).tag(config=True)
+
+    on_cell_error: t.Optional[t.Callable] = Any(
+        default_value=None,
+        allow_none=True,
+        help=dedent("""
+        A callable which executes when a cell execution results in an error.
+        This is executed even if errors are suppressed with `cell_allows_errors`.
+        Called with kwargs `cell`, and `cell_index`.
+        """),
+    ).tag(config=True)
+
     @default('kernel_manager_class')
     def _kernel_manager_class_default(self) -> KernelManager:
         """Use a dynamic default to avoid importing jupyter_client at startup"""
@@ -442,6 +481,7 @@ async def async_start_new_kernel_client(self) -> KernelClient:
             await self._async_cleanup_kernel()
             raise
         self.kc.allow_stdin = False
+        run_hook(sself.on_execution_start)
         return self.kc
 
     start_new_kernel_client = run_sync(async_start_new_kernel_client)
@@ -745,7 +785,11 @@ def _passed_deadline(self, deadline: int) -> bool:
             return True
         return False
 
-    def _check_raise_for_error(self, cell: NotebookNode, exec_reply: t.Optional[t.Dict]) -> None:
+    def _check_raise_for_error(
+            self,
+            cell: NotebookNode,
+            cell_index: int,
+            exec_reply: t.Optional[t.Dict]) -> None:
 
         if exec_reply is None:
             return None
@@ -760,8 +804,10 @@ def _check_raise_for_error(self, cell: NotebookNode, exec_reply: t.Optional[t.Di
             or "raises-exception" in cell.metadata.get("tags", [])
         )
 
-        if not cell_allows_errors:
-            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, cell_index=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,
@@ -821,11 +867,13 @@ async def async_execute_cell(
             self.allow_errors or "raises-exception" in cell.metadata.get("tags", [])
         )
 
+        run_hook(self.on_cell_start, cell=cell, cell_index=cell_index)
         parent_msg_id = await ensure_async(
             self.kc.execute(
                 cell.source, store_history=store_history, stop_on_error=not cell_allows_errors
             )
         )
+        run_hook(self.on_cell_complete, cell=cell, cell_index=cell_index)
         # We launched a code cell to execute
         self.code_cells_executed += 1
         exec_timeout = self._get_timeout(cell)
@@ -859,7 +907,7 @@ async def async_execute_cell(
 
         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

@@ -6,7 +6,8 @@
 import asyncio
 import inspect
 import sys
-from typing import Any, Awaitable, Callable, Union
+from typing import Any, Awaitable, Callable, Optional, Union
+from functools import partial
 
 
 def check_ipython() -> None:
@@ -102,3 +103,15 @@ async def ensure_async(obj: Union[Awaitable, Any]) -> Any:
         return result
     # obj doesn't need to be awaited
     return obj
+
+
+def run_hook(hook: Optional[Callable], **kwargs) -> None:
+    if hook is None:
+        return
+    if inspect.iscoroutinefunction(hook):
+        future = hook(**kwargs)
+    else:
+        loop = asyncio.get_event_loop()
+        hook_with_kwargs = partial(hook, **kwargs)
+        future = loop.run_in_executor(None, hook_with_kwargs)
+    asyncio.ensure_future(future)