Add on_cell_executed hook

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()