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

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 notebook or cell execution
+--------------------------------------------------
+There are several configurable hooks that allow the user to execute code before and
+after a notebook or a cell is executed. Each one is configured with a function that will be called in its
+respective place in the 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 `kernel_id`.
- """),
+ 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_cell_start: t.Optional[t.Callable] = Any(
+ on_notebook_complete: 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(
+ """
+ Called after kernel is cleaned up.
+ Called with kwargs `notebook`.
+ """
+ ),
  ).tag(config=True)
 
- on_cell_complete: t.Optional[t.Callable] = Any(
+ on_notebook_error: 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(
+ """
+ Called when the notebook encountered an error.
+ Called with kwargs `notebook`.
+ """
+ ),
  ).tag(config=True)
 
- on_cell_error: t.Optional[t.Callable] = Any(
+ on_cell_start: 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 before a cell is executed and before non-executing cells
+ are skipped.
+ Called with kwargs `cell` and `cell_index`.
+ """
+ ),
+ ).tag(config=True)
+
+ on_cell_execute: t.Optional[t.Callable] = Callable(
+ default_value=None,
+ allow_none=True,
+ 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] = 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`.
+ """
+ ),
+ ).tag(config=True)
+
+ 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`.
+ """
+ ),
  ).tag(config=True)
 
  @default('kernel_manager_class')
@@ -481,7 +534,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)
+ 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 +606,13 @@ def on_signal():
  await self.async_start_new_kernel_client()
  try:
  yield
+ except RuntimeError as e:
+ await run_hook(self.on_notebook_error, notebook=self.nb)
+ raise e
  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 +841,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 +857,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 +902,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 +922,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 +962,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