Addressing PR comments

Author: Golf Player

.bumpversion.cfg

@@ -227,33 +227,43 @@ class NotebookClient(LoggingConfigurable):
 
     kernel_manager_class: KernelManager = Type(config=True, help='The kernel manager class to use.')
 
-    on_kernel_create = Any(
+    on_execution_start = Any(
         default_value=None,
         allow_none=True,
-        help="""A callable which executes when the kernel is created.""",
+        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 = Any(
         default_value=None,
         allow_none=True,
-        help="""A callable which executes before a cell is executed.""",
+        help=dedent("""
+        A callable which executes before a cell is executed.
+        Called with kwargs `cell`, and `cell_index`.
+        """),
     ).tag(config=True)
 
     on_cell_complete = 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.
-            """
-        ),
+        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 = Any(
         default_value=None,
         allow_none=True,
-        help="""A callable which executes when a cell execution results in an error.""",
+        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')
@@ -418,7 +428,6 @@ async def async_start_new_kernel_client(self, **kwargs) -> t.Tuple[KernelClient,
 
         kernel_id = await ensure_async(self.km.start_kernel(extra_arguments=self.extra_arguments,
                                        **kwargs))
-        run_hook(self.on_kernel_create, kernel_id)
 
         # if self.km is not a KernelManager, it's probably a MultiKernelManager
         try:
@@ -442,6 +451,7 @@ async def async_start_new_kernel_client(self, **kwargs) -> t.Tuple[KernelClient,
             await self._async_cleanup_kernel()
             raise
         self.kc.allow_stdin = False
+        run_hook(self.on_execution_start, kernel_id=kernel_id)
         return self.kc, kernel_id
 
     start_new_kernel_client = run_sync(async_start_new_kernel_client)
@@ -740,7 +750,7 @@ def _check_raise_for_error(
         )
 
         if (exec_reply is not None) and exec_reply['content']['status'] == 'error':
-            run_hook(self.on_cell_error, cell, cell_index)
+            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'])
 
@@ -792,14 +802,15 @@ async def async_execute_cell(
             cell['metadata']['execution'] = {}
 
         self.log.debug("Executing cell:\n%s", cell.source)
-        run_hook(self.on_cell_start, cell, cell_index)
+        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 self.allow_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)

nbclient/client.py

@@ -7,6 +7,7 @@
 import sys
 import inspect
 from typing import Callable, Awaitable, Any, Union
+from functools import partial
 
 
 def check_ipython() -> None:
@@ -93,12 +94,13 @@ async def ensure_async(obj: Union[Awaitable, Any]) -> Any:
     return obj
 
 
-def run_hook(hook, *args):
+def run_hook(hook: Callable, **kwargs) -> None:
     if hook is None:
         return
     if inspect.iscoroutinefunction(hook):
-        future = hook(*args)
+        future = hook(**kwargs)
     else:
         loop = asyncio.get_event_loop()
-        future = loop.run_in_executor(None, hook, *args)
+        hook_with_kwargs = partial(hook, **kwargs)
+        future = loop.run_in_executor(None, hook_with_kwargs)
     asyncio.ensure_future(future)