Merge pull request #44 from davidbrochart/reset_kc

Add reset_kc option to reset_execution_trackers

Author: MSeal

nbclient/client.py

@@ -303,7 +303,6 @@ def __init__(self, nb, km=None, **kw):
     def reset_execution_trackers(self):
         """Resets any per-execution trackers.
         """
-        self.kc = None
         self.code_cells_executed = 0
         self._display_id_map = {}
         self.widget_state = {}
@@ -412,11 +411,12 @@ async def async_setup_kernel(self, **kwargs):
         """
         Context manager for setting up the kernel to execute a notebook.
 
-        The assigns the Kernel Manager (`self.km`) if missing and Kernel Client(`self.kc`).
+        This assigns the Kernel Manager (`self.km`) if missing and Kernel Client(`self.kc`).
 
         When control returns from the yield it stops the client's zmq channels, and shuts
         down the kernel.
         """
+        reset_kc = kwargs.pop('reset_kc', False)
         if self.km is None:
             self.start_kernel_manager()
 
@@ -425,17 +425,31 @@ async def async_setup_kernel(self, **kwargs):
         try:
             yield
         finally:
-            await self._async_cleanup_kernel()
+            if reset_kc:
+                await self._async_cleanup_kernel()
 
     async def async_execute(self, **kwargs):
         """
         Executes each code cell.
 
+        Parameters
+        ----------
+        kwargs :
+            Any option for `self.kernel_manager_class.start_kernel()`. Because
+            that defaults to AsyncKernelManager, this will likely include options
+            accepted by `AsyncKernelManager.start_kernel()``, which includes `cwd`.
+            If present, `reset_kc` is passed to `self.async_setup_kernel`:
+            if True, the kernel client will be reset and a new one will be created
+            and cleaned up after execution (default: False).
+
         Returns
         -------
         nb : NotebookNode
             The executed notebook.
         """
+        reset_kc = kwargs.get('reset_kc', False)
+        if reset_kc:
+            await self._async_cleanup_kernel()
         self.reset_execution_trackers()
 
         async with self.async_setup_kernel(**kwargs):

nbclient/tests/test_client.py

@@ -544,6 +544,29 @@ def test_force_raise_errors(self):
             else:
                 assert u"# üñîçø∂é".encode('utf8', 'replace') in str(exc.value)
 
+    def test_reset_kernel_client(self):
+        filename = os.path.join(current_dir, 'files', 'HelloWorld.ipynb')
+
+        with io.open(filename) as f:
+            input_nb = nbformat.read(f, 4)
+
+        executor = NotebookClient(
+            input_nb,
+            resources=self.build_resources(),
+        )
+
+        executor.execute()
+        # we didn't ask to reset the kernel client, a new one must have been created
+        kc = executor.kc
+        assert kc is not None
+        executor.execute()
+        # we didn't ask to reset the kernel client, the previously created one must have been reused
+        assert kc == executor.kc
+        executor.execute(reset_kc=True)
+        # we asked to reset the kernel client, the previous one must have been cleaned up,
+        # a new one must have been created and also cleaned up
+        assert executor.kc is None
+
     def test_custom_kernel_manager(self):
         from .fake_kernelmanager import FakeCustomKernelManager