Add on_cell_executed hook (#222)

* Add on_cell_executed hook

* Update doc and refactor tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

docs/client.rst

@@ -113,7 +113,7 @@ Here is the available hooks:
 - ``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-level hooks**: These hooks are called with at least 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.
@@ -123,8 +123,11 @@ 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_executed`` will run right after the code cell is executed.
 - ``on_cell_error`` will run if there is an error during cell execution.
 
+``on_cell_executed`` and ``on_cell_error`` are called with an extra parameter ``execute_reply=dict``.
+
 
 Handling errors and exceptions
 ------------------------------

nbclient/client.py

@@ -340,14 +340,26 @@ class NotebookClient(LoggingConfigurable):
  ),
  ).tag(config=True)
 
+ on_cell_executed: t.Optional[t.Callable] = Callable(
+ default_value=None,
+ allow_none=True,
+ help=dedent(
+ """
+ A callable which executes just after a code cell is executed, whether
+ or not it results in an error.
+ Called with kwargs ``cell``, ``cell_index`` and ``execute_reply``.
+ """
+ ),
+ ).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``.
+ Called with kwargs ``cell`, ``cell_index`` and ``execute_reply``.
  """
  ),
  ).tag(config=True)
@@ -857,7 +869,9 @@ async def _check_raise_for_error(
  or exec_reply_content.get('ename') in self.allow_error_names
  or "raises-exception" in cell.metadata.get("tags", [])
  )
- await run_hook(self.on_cell_error, cell=cell, cell_index=cell_index)
+ await run_hook(
+ self.on_cell_error, cell=cell, cell_index=cell_index, execute_reply=exec_reply
+ )
  if not cell_allows_errors:
  raise CellExecutionError.from_cell_and_msg(cell, exec_reply_content)
 
@@ -962,6 +976,9 @@ async def async_execute_cell(
 
  if execution_count:
  cell['execution_count'] = execution_count
+ await run_hook(
+ self.on_cell_executed, cell=cell, cell_index=cell_index, execute_reply=exec_reply
+ )
  await self._check_raise_for_error(cell, cell_index, exec_reply)
  self.nb['cells'][cell_index] = cell
  return cell

nbclient/tests/test_client.py

@@ -40,13 +40,40 @@
  "on_cell_start",
  "on_cell_execute",
  "on_cell_complete",
+ "on_cell_executed",
  "on_cell_error",
  "on_notebook_start",
  "on_notebook_complete",
  "on_notebook_error",
 ]
 
 
+def get_executor_with_hooks(nb=None, executor=None, async_hooks=False):
+ if async_hooks:
+ hooks = {key: AsyncMock() for key in hook_methods}
+ else:
+ hooks = {key: MagicMock() for key in hook_methods}
+ if nb is not None:
+ if executor is not None:
+ raise RuntimeError("Cannot pass nb and executor at the same time")
+ executor = NotebookClient(nb)
+ for k, v in hooks.items():
+ setattr(executor, k, v)
+ return executor, hooks
+
+
+EXECUTE_REPLY_OK = {
+ 'parent_header': {'msg_id': 'fake_id'},
+ 'content': {'status': 'ok', 'execution_count': 1},
+}
+EXECUTE_REPLY_ERROR = {
+ 'parent_header': {'msg_id': 'fake_id'},
+ 'content': {'status': 'error'},
+ 'msg_type': 'execute_reply',
+ 'header': {'msg_type': 'execute_reply'},
+}
+
+
 class AsyncMock(Mock):
  pass
 
@@ -760,77 +787,79 @@ def test_execution_hook(self):
  filename = os.path.join(current_dir, 'files', 'HelloWorld.ipynb')
  with open(filename) as f:
  input_nb = nbformat.read(f, 4)
- hooks = [MagicMock() for i in range(7)]
- executor = NotebookClient(input_nb)
- for executor_hook, hook in zip(hook_methods, hooks):
- setattr(executor, executor_hook, hook)
+ executor, hooks = get_executor_with_hooks(nb=input_nb)
  executor.execute()
- for hook in hooks[:3]:
- hook.assert_called_once()
- hooks[3].assert_not_called()
- for hook in hooks[4:6]:
- hook.assert_called_once()
- hooks[6].assert_not_called()
+ hooks["on_cell_start"].assert_called_once()
+ hooks["on_cell_execute"].assert_called_once()
+ hooks["on_cell_complete"].assert_called_once()
+ hooks["on_cell_executed"].assert_called_once()
+ hooks["on_cell_error"].assert_not_called()
+ hooks["on_notebook_start"].assert_called_once()
+ hooks["on_notebook_complete"].assert_called_once()
+ hooks["on_notebook_error"].assert_not_called()
 
  def test_error_execution_hook_error(self):
  filename = os.path.join(current_dir, 'files', 'Error.ipynb')
  with open(filename) as f:
  input_nb = nbformat.read(f, 4)
- hooks = [MagicMock() for i in range(7)]
- executor = NotebookClient(input_nb)
- for executor_hook, hook in zip(hook_methods, hooks):
- setattr(executor, executor_hook, hook)
+ executor, hooks = get_executor_with_hooks(nb=input_nb)
  with pytest.raises(CellExecutionError):
  executor.execute()
- for hook in hooks[:5]:
- hook.assert_called_once()
- hooks[6].assert_not_called()
+ hooks["on_cell_start"].assert_called_once()
+ hooks["on_cell_execute"].assert_called_once()
+ hooks["on_cell_complete"].assert_called_once()
+ hooks["on_cell_executed"].assert_called_once()
+ hooks["on_cell_error"].assert_called_once()
+ hooks["on_notebook_start"].assert_called_once()
+ hooks["on_notebook_complete"].assert_called_once()
+ hooks["on_notebook_error"].assert_not_called()
 
  def test_error_notebook_hook(self):
  filename = os.path.join(current_dir, 'files', 'Autokill.ipynb')
  with open(filename) as f:
  input_nb = nbformat.read(f, 4)
- hooks = [MagicMock() for i in range(7)]
- executor = NotebookClient(input_nb)
- for executor_hook, hook in zip(hook_methods, hooks):
- setattr(executor, executor_hook, hook)
+ executor, hooks = get_executor_with_hooks(nb=input_nb)
  with pytest.raises(RuntimeError):
  executor.execute()
- for hook in hooks[:3]:
- hook.assert_called_once()
- hooks[3].assert_not_called()
- for hook in hooks[4:]:
- hook.assert_called_once()
+ hooks["on_cell_start"].assert_called_once()
+ hooks["on_cell_execute"].assert_called_once()
+ hooks["on_cell_complete"].assert_called_once()
+ hooks["on_cell_executed"].assert_not_called()
+ hooks["on_cell_error"].assert_not_called()
+ hooks["on_notebook_start"].assert_called_once()
+ hooks["on_notebook_complete"].assert_called_once()
+ hooks["on_notebook_error"].assert_called_once()
 
  def test_async_execution_hook(self):
  filename = os.path.join(current_dir, 'files', 'HelloWorld.ipynb')
  with open(filename) as f:
  input_nb = nbformat.read(f, 4)
- hooks = [AsyncMock() for i in range(7)]
- executor = NotebookClient(input_nb)
- for executor_hook, hook in zip(hook_methods, hooks):
- setattr(executor, executor_hook, hook)
+ executor, hooks = get_executor_with_hooks(nb=input_nb)
  executor.execute()
- for hook in hooks[:3]:
- hook.assert_called_once()
- hooks[3].assert_not_called()
- for hook in hooks[4:6]:
- hook.assert_called_once()
- hooks[6].assert_not_called()
+ hooks["on_cell_start"].assert_called_once()
+ hooks["on_cell_execute"].assert_called_once()
+ hooks["on_cell_complete"].assert_called_once()
+ hooks["on_cell_executed"].assert_called_once()
+ hooks["on_cell_error"].assert_not_called()
+ hooks["on_notebook_start"].assert_called_once()
+ hooks["on_notebook_complete"].assert_called_once()
+ hooks["on_notebook_error"].assert_not_called()
 
  def test_error_async_execution_hook(self):
  filename = os.path.join(current_dir, 'files', 'Error.ipynb')
  with open(filename) as f:
  input_nb = nbformat.read(f, 4)
- hooks = [AsyncMock() for i in range(7)]
- executor = NotebookClient(input_nb)
- for executor_hook, hook in zip(hook_methods, hooks):
- setattr(executor, executor_hook, hook)
+ executor, hooks = get_executor_with_hooks(nb=input_nb)
  with pytest.raises(CellExecutionError):
- executor.execute().execute()
- for hook in hooks[:5]:
- hook.assert_called_once()
- hooks[6].assert_not_called()
+ executor.execute()
+ hooks["on_cell_start"].assert_called_once()
+ hooks["on_cell_execute"].assert_called_once()
+ hooks["on_cell_complete"].assert_called_once()
+ hooks["on_cell_executed"].assert_called_once()
+ hooks["on_cell_error"].assert_called_once()
+ hooks["on_notebook_start"].assert_called_once()
+ hooks["on_notebook_complete"].assert_called_once()
+ hooks["on_notebook_error"].assert_not_called()
 
 
 class TestRunCell(NBClientTestsBase):
@@ -1618,14 +1647,18 @@ def test_no_source(self, executor, cell_mock, message_mock):
 
  @prepare_cell_mocks()
  def test_cell_hooks(self, executor, cell_mock, message_mock):
- hooks = [MagicMock() for i in range(7)]
- for executor_hook, hook in zip(hook_methods, hooks):
- setattr(executor, executor_hook, hook)
+ executor, hooks = get_executor_with_hooks(executor=executor)
  executor.execute_cell(cell_mock, 0)
- for hook in hooks[:3]:
- hook.assert_called_once_with(cell=cell_mock, cell_index=0)
- for hook in hooks[4:]:
- hook.assert_not_called()
+ hooks["on_cell_start"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_execute"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_complete"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_executed"].assert_called_once_with(
+ cell=cell_mock, cell_index=0, execute_reply=EXECUTE_REPLY_OK
+ )
+ hooks["on_cell_error"].assert_not_called()
+ hooks["on_notebook_start"].assert_not_called()
+ hooks["on_notebook_complete"].assert_not_called()
+ hooks["on_notebook_error"].assert_not_called()
 
  @prepare_cell_mocks(
  {
@@ -1641,15 +1674,21 @@ def test_cell_hooks(self, executor, cell_mock, message_mock):
  },
  )
  def test_error_cell_hooks(self, executor, cell_mock, message_mock):
- hooks = [MagicMock() for i in range(7)]
- for executor_hook, hook in zip(hook_methods, hooks):
- setattr(executor, executor_hook, hook)
+ executor, hooks = get_executor_with_hooks(executor=executor)
  with self.assertRaises(CellExecutionError):
  executor.execute_cell(cell_mock, 0)
- for hook in hooks[:4]:
- hook.assert_called_once_with(cell=cell_mock, cell_index=0)
- for hook in hooks[5:]:
- hook.assert_not_called()
+ hooks["on_cell_start"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_execute"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_complete"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_executed"].assert_called_once_with(
+ cell=cell_mock, cell_index=0, execute_reply=EXECUTE_REPLY_ERROR
+ )
+ hooks["on_cell_error"].assert_called_once_with(
+ cell=cell_mock, cell_index=0, execute_reply=EXECUTE_REPLY_ERROR
+ )
+ hooks["on_notebook_start"].assert_not_called()
+ hooks["on_notebook_complete"].assert_not_called()
+ hooks["on_notebook_error"].assert_not_called()
 
  @prepare_cell_mocks(
  reply_msg={
@@ -1661,25 +1700,31 @@ def test_error_cell_hooks(self, executor, cell_mock, message_mock):
  )
  def test_non_code_cell_hooks(self, executor, cell_mock, message_mock):
  cell_mock = NotebookNode(source='"foo" = "bar"', metadata={}, cell_type='raw', outputs=[])
- hooks = [MagicMock() for i in range(7)]
- for executor_hook, hook in zip(hook_methods, hooks):
- setattr(executor, executor_hook, hook)
+ executor, hooks = get_executor_with_hooks(executor=executor)
  executor.execute_cell(cell_mock, 0)
- for hook in hooks[:1]:
- hook.assert_called_once_with(cell=cell_mock, cell_index=0)
- for hook in hooks[1:]:
- hook.assert_not_called()
+ hooks["on_cell_start"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_execute"].assert_not_called()
+ hooks["on_cell_complete"].assert_not_called()
+ hooks["on_cell_executed"].assert_not_called()
+ hooks["on_cell_error"].assert_not_called()
+ hooks["on_notebook_start"].assert_not_called()
+ hooks["on_notebook_complete"].assert_not_called()
+ hooks["on_notebook_error"].assert_not_called()
 
  @prepare_cell_mocks()
  def test_async_cell_hooks(self, executor, cell_mock, message_mock):
- hooks = [AsyncMock() for i in range(7)]
- for executor_hook, hook in zip(hook_methods, hooks):
- setattr(executor, executor_hook, hook)
+ executor, hooks = get_executor_with_hooks(executor=executor, async_hooks=True)
  executor.execute_cell(cell_mock, 0)
- for hook in hooks[:3]:
- hook.assert_called_once_with(cell=cell_mock, cell_index=0)
- for hook in hooks[4:]:
- hook.assert_not_called()
+ hooks["on_cell_start"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_execute"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_complete"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_executed"].assert_called_once_with(
+ cell=cell_mock, cell_index=0, execute_reply=EXECUTE_REPLY_OK
+ )
+ hooks["on_cell_error"].assert_not_called()
+ hooks["on_notebook_start"].assert_not_called()
+ hooks["on_notebook_complete"].assert_not_called()
+ hooks["on_notebook_error"].assert_not_called()
 
  @prepare_cell_mocks(
  {
@@ -1695,12 +1740,18 @@ def test_async_cell_hooks(self, executor, cell_mock, message_mock):
  },
  )
  def test_error_async_cell_hooks(self, executor, cell_mock, message_mock):
- hooks = [AsyncMock() for i in range(7)]
- for executor_hook, hook in zip(hook_methods, hooks):
- setattr(executor, executor_hook, hook)
+ executor, hooks = get_executor_with_hooks(executor=executor, async_hooks=True)
  with self.assertRaises(CellExecutionError):
  executor.execute_cell(cell_mock, 0)
- for hook in hooks[:4]:
- hook.assert_called_once_with(cell=cell_mock, cell_index=0)
- for hook in hooks[4:]:
- hook.assert_not_called()
+ hooks["on_cell_start"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_execute"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_complete"].assert_called_once_with(cell=cell_mock, cell_index=0)
+ hooks["on_cell_executed"].assert_called_once_with(
+ cell=cell_mock, cell_index=0, execute_reply=EXECUTE_REPLY_ERROR
+ )
+ hooks["on_cell_error"].assert_called_once_with(
+ cell=cell_mock, cell_index=0, execute_reply=EXECUTE_REPLY_ERROR
+ )
+ hooks["on_notebook_start"].assert_not_called()
+ hooks["on_notebook_complete"].assert_not_called()
+ hooks["on_notebook_error"].assert_not_called()