Forward cleanup_kc through execute

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)

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"""