diff --git a/nbclient/client.py b/nbclient/client.py index c377063c..1bc46243 100644 --- a/nbclient/client.py +++ b/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): diff --git a/nbclient/tests/test_client.py b/nbclient/tests/test_client.py index 37ef3958..9782f6cb 100644 --- a/nbclient/tests/test_client.py +++ b/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