From 75bc0e742bb0ec570fee8027f7fbcb6c47019e06 Mon Sep 17 00:00:00 2001 From: David Brochart Date: Thu, 9 Apr 2020 23:04:00 +0200 Subject: [PATCH 1/3] Add reset_kc option to reset_execution_trackers, async_execute and execute methods (defaults to False) --- nbclient/client.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/nbclient/client.py b/nbclient/client.py index c377063c..7383b6a7 100644 --- a/nbclient/client.py +++ b/nbclient/client.py @@ -298,12 +298,13 @@ def __init__(self, nb, km=None, **kw): super().__init__(**kw) self.nb = nb self.km = km - self.reset_execution_trackers() + self.reset_execution_trackers(reset_kc=True) - def reset_execution_trackers(self): + def reset_execution_trackers(self, reset_kc=False): """Resets any per-execution trackers. """ - self.kc = None + if reset_kc: + self.kc = None self.code_cells_executed = 0 self._display_id_map = {} self.widget_state = {} @@ -412,7 +413,7 @@ 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. @@ -431,12 +432,23 @@ 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.reset_execution_trackers`: + if True, the kernel client will be reset and a new one will be created + (default: False). + Returns ------- nb : NotebookNode The executed notebook. """ - self.reset_execution_trackers() + reset_kc = kwargs.pop('reset_kc', False) + self.reset_execution_trackers(reset_kc=reset_kc) async with self.async_setup_kernel(**kwargs): self.log.info("Executing notebook with kernel: %s" % self.kernel_name) From dc22d8e78fbff28da52a570496409e5b83057992 Mon Sep 17 00:00:00 2001 From: David Brochart Date: Fri, 10 Apr 2020 09:33:34 +0200 Subject: [PATCH 2/3] Add test for kernel client reset --- nbclient/client.py | 20 +++++++++++--------- nbclient/tests/test_client.py | 23 +++++++++++++++++++++++ 2 files changed, 34 insertions(+), 9 deletions(-) diff --git a/nbclient/client.py b/nbclient/client.py index 7383b6a7..1bc46243 100644 --- a/nbclient/client.py +++ b/nbclient/client.py @@ -298,13 +298,11 @@ def __init__(self, nb, km=None, **kw): super().__init__(**kw) self.nb = nb self.km = km - self.reset_execution_trackers(reset_kc=True) + self.reset_execution_trackers() - def reset_execution_trackers(self, reset_kc=False): + def reset_execution_trackers(self): """Resets any per-execution trackers. """ - if reset_kc: - self.kc = None self.code_cells_executed = 0 self._display_id_map = {} self.widget_state = {} @@ -418,6 +416,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) if self.km is None: self.start_kernel_manager() @@ -426,7 +425,8 @@ 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): """ @@ -438,17 +438,19 @@ 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.reset_execution_trackers`: + 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 - (default: False). + and cleaned up after execution (default: False). Returns ------- nb : NotebookNode The executed notebook. """ - reset_kc = kwargs.pop('reset_kc', False) - self.reset_execution_trackers(reset_kc=reset_kc) + 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): self.log.info("Executing notebook with kernel: %s" % self.kernel_name) diff --git a/nbclient/tests/test_client.py b/nbclient/tests/test_client.py index 37ef3958..10e40f11 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 + self.assertNotEqual(kc, None) + executor.execute() + # we didn't ask to reset the kernel client, the previously created one must have been reused + self.assertEqual(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 + self.assertEqual(executor.kc, None) + def test_custom_kernel_manager(self): from .fake_kernelmanager import FakeCustomKernelManager From 7d3a1c3e44c72d4e7defae81efdd95ef1b40cb1b Mon Sep 17 00:00:00 2001 From: David Brochart Date: Fri, 10 Apr 2020 13:29:02 +0200 Subject: [PATCH 3/3] Use python's assert --- nbclient/tests/test_client.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nbclient/tests/test_client.py b/nbclient/tests/test_client.py index 10e40f11..9782f6cb 100644 --- a/nbclient/tests/test_client.py +++ b/nbclient/tests/test_client.py @@ -558,14 +558,14 @@ def test_reset_kernel_client(self): executor.execute() # we didn't ask to reset the kernel client, a new one must have been created kc = executor.kc - self.assertNotEqual(kc, None) + assert kc is not None executor.execute() # we didn't ask to reset the kernel client, the previously created one must have been reused - self.assertEqual(kc, executor.kc) + 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 - self.assertEqual(executor.kc, None) + assert executor.kc is None def test_custom_kernel_manager(self): from .fake_kernelmanager import FakeCustomKernelManager