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