Merge pull request #53 from danielfrg/patch-1

Adding reset_kc to setup_kernel

nbclient/client.py

@@ -414,6 +414,8 @@ def setup_kernel(self, **kwargs):
         When control returns from the yield it stops the client's zmq channels, and shuts
         down the kernel.
         """
+        cleanup_kc = kwargs.pop('cleanup_kc', True)
+
         # Can't use run_until_complete on an asynccontextmanager function :(
         if self.km is None:
             self.start_kernel_manager()
@@ -423,7 +425,8 @@ def setup_kernel(self, **kwargs):
         try:
             yield
         finally:
-            self._cleanup_kernel()
+            if cleanup_kc:
+                self._cleanup_kernel()
 
     @asynccontextmanager
     async def async_setup_kernel(self, **kwargs):
@@ -435,7 +438,7 @@ async def async_setup_kernel(self, **kwargs):
         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)
+        cleanup_kc = kwargs.pop('cleanup_kc', True)
         if self.km is None:
             self.start_kernel_manager()
 
@@ -444,10 +447,10 @@ async def async_setup_kernel(self, **kwargs):
         try:
             yield
         finally:
-            if reset_kc:
+            if cleanup_kc:
                 await self._async_cleanup_kernel()
 
-    async def async_execute(self, **kwargs):
+    async def async_execute(self, reset_kc=False, **kwargs):
         """
         Executes each code cell.
 
@@ -457,16 +460,15 @@ async def async_execute(self, **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).
+
+            `reset_kc` if True, the kernel client will be reset and a new one
+            will be created (default: False).
 
         Returns
         -------
         nb : NotebookNode
             The executed notebook.
         """
-        reset_kc = kwargs.get('reset_kc', False)
         if reset_kc and self.km:
             await self._async_cleanup_kernel()
         self.reset_execution_trackers()

nbclient/tests/test_client.py

@@ -594,18 +594,40 @@ def test_reset_kernel_client(self):
             resources=self.build_resources(),
         )
 
-        executor.execute()
+        executor.execute(cleanup_kc=False)
         # 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()
+
+        executor.execute(cleanup_kc=False)
         # 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)
+
+        executor.execute(reset_kc=True, cleanup_kc=False)
         # 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
+        # a new one must have been created
+        assert kc != executor.kc
+
+    def test_cleanup_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 asked to cleanup the kernel client (default is True)
         assert executor.kc is None
 
+        executor.execute(cleanup_kc=False)
+        # we didn't ask to reset the kernel client
+        # a new one must have been created and should still be available
+        assert executor.kc is not None
+
     def test_custom_kernel_manager(self):
         from .fake_kernelmanager import FakeCustomKernelManager