diff --git a/nbclient/client.py b/nbclient/client.py index f60361c..3ed9962 100644 --- a/nbclient/client.py +++ b/nbclient/client.py @@ -1299,6 +1299,7 @@ def execute( nb: NotebookNode, cwd: str | None = None, km: KernelManager | None = None, + cleanup_kc: bool | None = None, **kwargs: t.Any, ) -> NotebookNode: """Execute a notebook's code, updating outputs within the notebook object. @@ -1314,10 +1315,16 @@ def execute( If supplied, the kernel will run in this directory km : AsyncKernelManager, optional If supplied, the specified kernel manager will be used for code execution. + cleanup_kc : bool, optional + Whether to stop the kernel client's channels and shut down the kernel after execution. + If ``None`` or omitted, ``NotebookClient.execute`` retains its default behavior. kwargs : Any other options for NotebookClient, e.g. timeout, kernel_name """ resources = {} if cwd is not None: resources["metadata"] = {"path": cwd} - return NotebookClient(nb=nb, resources=resources, km=km, **kwargs).execute() + client = NotebookClient(nb=nb, resources=resources, km=km, **kwargs) + if cleanup_kc is None: + return client.execute() + return client.execute(cleanup_kc=cleanup_kc) diff --git a/tests/test_client.py b/tests/test_client.py index fbc9ce3..b776ee9 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -588,6 +588,28 @@ def stop_channels(self) -> None: assert not km.has_kernel # type:ignore[unreachable] +@pytest.mark.parametrize( + ("kwargs", "expected_kwargs"), + [ + ({}, {}), + ({"cleanup_kc": None}, {}), + ({"cleanup_kc": True}, {"cleanup_kc": True}), + ({"cleanup_kc": False}, {"cleanup_kc": False}), + ], + ids=["default", "none", "true", "false"], +) +def test_execute_function_cleanup_kc(monkeypatch, kwargs, expected_kwargs): + notebook = nbformat.v4.new_notebook() + executed_notebook = nbformat.v4.new_notebook() + execute_mock = MagicMock(return_value=executed_notebook) + monkeypatch.setattr(NotebookClient, "execute", execute_mock) + + result = execute(notebook, km=MagicMock(spec=KernelManager), **kwargs) + + execute_mock.assert_called_once_with(**expected_kwargs) + assert result is executed_notebook + + class TestExecute(NBClientTestsBase): """Contains test functions for execute.py"""