Rebased with master and added tests

Run_hook is now async and renamed util to test_util so it gets picked up by pytest.
Also added new hooks: on_notebook_error, on_cell_execution
Updated docs

Author: devintang3

docs/client.rst

@@ -96,6 +96,36 @@ on both versions. Here the traitlet ``kernel_name`` helps simplify and
 maintain consistency: we can just run a notebook twice, specifying first
 "python2" and then "python3" as the kernel name.
 
+Hooks before and after cell execution
+-------------------------------------
+There are several configurable hooks that allow the user to execute code before and
+after a cell is executed. Each one is configured with a function that will be called in its
+respective place in the cell execution pipeline.
+Each is described below:
+
+**Notebook-level hooks**: These hooks are called with a single extra parameter:
+
+- ``notebook=NotebookNode``: the current notebook being executed.
+
+Here is the available hooks:
+
+- ``on_notebook_start`` will run when the notebook client is initialized, before any execution has happened.
+- ``on_notebook_complete`` will run when the notebook client has finished executing, after kernel cleanup.
+- ``on_notebook_error`` will run when the notebook client has encountered an exception before kernel cleanup.
+
+**Cell-level hooks**: These hooks are called with two parameters:
+
+- ``cell=NotebookNode``: a reference to the current cell.
+- ``cell_index=int``: the index of the cell in the current notebook's list of cells
+
+Here are the available hooks:
+
+- ``on_cell_start`` will run for all cell types before the cell is executed.
+- ``on_cell_execute`` will run right before the code cell is executed.
+- ``on_cell_complete`` will run after execution, if the cell is executed with no errors.
+- ``on_cell_error`` will run if there is an error during cell execution.
+
+
 Handling errors and exceptions
 ------------------------------
 

nbclient/client.py

@@ -15,7 +15,18 @@
 from jupyter_client.client import KernelClient
 from nbformat import NotebookNode
 from nbformat.v4 import output_from_msg
-from traitlets import Any, Bool, Dict, Enum, Integer, List, Type, Unicode, default
+from traitlets import (
+    Any,
+    Bool,
+    Callable,
+    Dict,
+    Enum,
+    Integer,
+    List,
+    Type,
+    Unicode,
+    default,
+)
 from traitlets.config.configurable import LoggingConfigurable
 
 from .exceptions import (
@@ -26,7 +37,7 @@
     DeadKernelError,
 )
 from .output_widget import OutputWidget
-from .util import ensure_async, run_sync, run_hook
+from .util import ensure_async, run_hook, run_sync
 
 
 def timestamp(msg: Optional[Dict] = None) -> str:
@@ -261,43 +272,85 @@ 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(
+    on_notebook_start: t.Optional[t.Callable] = Callable(
+        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 `notebook`
+            """
+        ),
+    ).tag(config=True)
+
+    on_notebook_complete: t.Optional[t.Callable] = Callable(
+        default_value=None,
+        allow_none=True,
+        help=dedent(
+            """
+            Called after kernel is cleaned up.
+            Called with kwargs `notebook`
+            """
+        ),
+    ).tag(config=True)
+
+    on_notebook_error: t.Optional[t.Callable] = Callable(
+        default_value=None,
+        allow_none=True,
+        help=dedent(
+            """
+            Called when the notebook encountered an error.
+            Called with kwargs `notebook`
+            """
+        ),
+    ).tag(config=True)
+
+    on_cell_start: t.Optional[t.Callable] = Callable(
         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`.
-        """),
+        help=dedent(
+            """
+            A callable which executes before a cell is executed and before non-executing cells
+            are skipped.
+            Called with kwargs `cell` and `cell_index`.
+            """
+        ),
     ).tag(config=True)
 
-    on_cell_start: t.Optional[t.Callable] = Any(
+    on_cell_execute: t.Optional[t.Callable] = Callable(
         default_value=None,
         allow_none=True,
-        help=dedent("""
-        A callable which executes before a cell is executed.
-        Called with kwargs `cell`, and `cell_index`.
-        """),
+        help=dedent(
+            """
+            A callable which executes just before a code cell is executed.
+            Called with kwargs `cell` and `cell_index`.
+            """
+        ),
     ).tag(config=True)
 
-    on_cell_complete: t.Optional[t.Callable] = Any(
+    on_cell_complete: t.Optional[t.Callable] = Callable(
         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`.
-        """),
+        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(
+    on_cell_error: t.Optional[t.Callable] = Callable(
         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`.
-        """),
+        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')
@@ -478,10 +531,11 @@ async def async_start_new_kernel_client(self) -> KernelClient:
         try:
             await ensure_async(self.kc.wait_for_ready(timeout=self.startup_timeout))
         except RuntimeError:
+            await run_hook(self.on_notebook_error, notebook=self.nb)
             await self._async_cleanup_kernel()
             raise
         self.kc.allow_stdin = False
-        run_hook(sself.on_execution_start)
+        await run_hook(self.on_notebook_start, notebook=self.nb)
         return self.kc
 
     start_new_kernel_client = run_sync(async_start_new_kernel_client)
@@ -553,10 +607,13 @@ def on_signal():
             await self.async_start_new_kernel_client()
         try:
             yield
+        except RuntimeError:
+            await run_hook(self.on_notebook_error, notebook=self.nb)
+            raise
         finally:
             if cleanup_kc:
                 await self._async_cleanup_kernel()
-
+            await run_hook(self.on_notebook_complete, notebook=self.nb)
             atexit.unregister(self._cleanup_kernel)
             try:
                 loop.remove_signal_handler(signal.SIGINT)
@@ -785,11 +842,9 @@ def _passed_deadline(self, deadline: int) -> bool:
             return True
         return False
 
-    def _check_raise_for_error(
-            self,
-            cell: NotebookNode,
-            cell_index: int,
-            exec_reply: t.Optional[t.Dict]) -> None:
+    async 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
@@ -803,11 +858,9 @@ def _check_raise_for_error(
             or exec_reply_content.get('ename') in self.allow_error_names
             or "raises-exception" in cell.metadata.get("tags", [])
         )
-
-        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'])
+        await run_hook(self.on_cell_error, cell=cell, cell_index=cell_index)
+        if not cell_allows_errors:
+            raise CellExecutionError.from_cell_and_msg(cell, exec_reply_content)
 
     async def async_execute_cell(
         self,
@@ -850,6 +903,9 @@ async def async_execute_cell(
             The cell which was just processed.
         """
         assert self.kc is not None
+
+        await run_hook(self.on_cell_start, cell=cell, cell_index=cell_index)
+
         if cell.cell_type != 'code' or not cell.source.strip():
             self.log.debug("Skipping non-executing cell %s", cell_index)
             return cell
@@ -867,13 +923,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)
+        await run_hook(self.on_cell_execute, 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)
+        await 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)
@@ -907,7 +963,7 @@ async def async_execute_cell(
 
         if execution_count:
             cell['execution_count'] = execution_count
-        self._check_raise_for_error(cell, cell_index, exec_reply)
+        await self._check_raise_for_error(cell, cell_index, exec_reply)
         self.nb['cells'][cell_index] = cell
         return cell