feat: allow passing kernel manager to parser

myst_nb/core/execute/__init__.py

@@ -10,6 +10,7 @@
 
 if TYPE_CHECKING:
     from nbformat import NotebookNode
+    from jupyter_client.manager import KernelManager
 
     from myst_nb.core.config import NbParserConfig
     from myst_nb.core.loggers import LoggerType
@@ -21,6 +22,8 @@ def create_client(
     nb_config: NbParserConfig,
     logger: LoggerType,
     read_fmt: None | dict = None,
+    *,
+    km: KernelManager | None = None,
 ) -> NotebookClientBase:
     """Create a notebook execution client, to update its outputs.
 
@@ -59,12 +62,14 @@ def create_client(
         return NotebookClientBase(notebook, path, nb_config, logger)
 
     if nb_config.execution_mode in ("auto", "force"):
-        return NotebookClientDirect(notebook, path, nb_config, logger)
+        return NotebookClientDirect(notebook, path, nb_config, logger, km=km)
 
     if nb_config.execution_mode == "cache":
-        return NotebookClientCache(notebook, path, nb_config, logger, read_fmt=read_fmt)
+        return NotebookClientCache(
+            notebook, path, nb_config, logger, read_fmt=read_fmt, km=km
+        )
 
     if nb_config.execution_mode == "inline":
-        return NotebookClientInline(notebook, path, nb_config, logger)
+        return NotebookClientInline(notebook, path, nb_config, logger, km=km)
 
     return NotebookClientBase(notebook, path, nb_config, logger)

myst_nb/core/execute/base.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any
+from typing import Any, TYPE_CHECKING
 
 from nbformat import NotebookNode
 from typing_extensions import TypedDict, final
@@ -13,6 +13,9 @@
 from myst_nb.core.nb_to_tokens import nb_node_to_dict
 from myst_nb.ext.glue import extract_glue_data
 
+if TYPE_CHECKING:
+    from jupyter_client import KernelManager
+
 
 class ExecutionResult(TypedDict):
     """Result of executing a notebook."""
@@ -54,13 +57,16 @@ def __init__(
         path: Path | None,
         nb_config: NbParserConfig,
         logger: LoggerType,
+        *,
+        km: KernelManager | None = None,
         **kwargs: Any,
     ):
         """Initialize the client."""
         self._notebook = notebook
         self._path = path
         self._nb_config = nb_config
         self._logger = logger
+        self._km = km
         self._kwargs = kwargs
 
         self._glue_data: dict[str, NotebookNode] = {}

myst_nb/core/execute/cache.py

@@ -76,6 +76,7 @@ def start_client(self):
                 allow_errors=self.nb_config.execution_allow_errors,
                 timeout=self.nb_config.execution_timeout,
                 meta_override=True,  # TODO still support this?
+                km=self._km,
             )
 
         # handle success / failure cases

myst_nb/core/execute/direct.py

@@ -44,6 +44,7 @@ def start_client(self):
                 allow_errors=self.nb_config.execution_allow_errors,
                 timeout=self.nb_config.execution_timeout,
                 meta_override=True,  # TODO still support this?
+                km=self._km,
             )
 
         if result.err is not None:

myst_nb/core/execute/inline.py

@@ -56,6 +56,7 @@ def start_client(self):
             resources=resources,
             allow_errors=self.nb_config.execution_allow_errors,
             timeout=self.nb_config.execution_timeout,
+            km=self._km,
         )
         self._client.reset_execution_trackers()
         if self._client.km is None:

myst_nb/docutils_.py

@@ -6,7 +6,7 @@
 from functools import lru_cache, partial
 from importlib import resources as import_resources
 import os
-from typing import Any
+from typing import Any, TYPE_CHECKING
 
 from docutils import nodes
 from docutils.core import default_description, publish_cmdline
@@ -46,6 +46,10 @@
 from myst_nb.ext.glue import load_glue_docutils
 from myst_nb.warnings_ import MystNBWarnings, create_warning
 
+if TYPE_CHECKING:
+    from jupyter_client import KernelManager
+
+
 DOCUTILS_EXCLUDED_ARGS = list(
     {f.name for f in NbParserConfig.get_fields() if f.metadata.get("docutils_exclude")}
 )
@@ -83,6 +87,10 @@ class Parser(MystParser):
 
     config_section = "myst-nb parser"
 
+    def __init__(self, *args, km: KernelManager | None = None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.km = km
+
     def parse(self, inputstring: str, document: nodes.document) -> None:
         # register/unregister special directives and roles
         app = get_nb_roles_directives()
@@ -187,7 +195,9 @@ def _parse(self, inputstring: str, document: nodes.document) -> None:
 
         # open the notebook execution client,
         # this may execute the notebook immediately or during the page render
-        with create_client(notebook, document_source, nb_config, logger) as nb_client:
+        with create_client(
+            notebook, document_source, nb_config, logger, km=self.km
+        ) as nb_client:
             mdit_parser.options["nb_client"] = nb_client
             # convert to docutils AST, which is added to the document
             mdit_renderer.render(mdit_tokens, mdit_parser.options, mdit_env)

myst_nb/sphinx_.py

@@ -7,7 +7,7 @@
 import json
 from pathlib import Path
 import re
-from typing import Any, DefaultDict, cast
+from typing import Any, DefaultDict, cast, TYPE_CHECKING
 
 from docutils import nodes
 from markdown_it.token import Token
@@ -41,6 +41,10 @@
 )
 from myst_nb.warnings_ import MystNBWarnings, create_warning
 
+if TYPE_CHECKING:
+    from jupyter_client import KernelManager
+
+
 SPHINX_LOGGER = sphinx_logging.getLogger(__name__)
 
 
@@ -64,6 +68,10 @@ class Parser(MystParser):
 
     env: SphinxEnvType
 
+    def __init__(self, *args, km: KernelManager | None = None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.km = km
+
     def parse(self, inputstring: str, document: nodes.document) -> None:
         """Parse source text.
 
@@ -154,7 +162,7 @@ def parse(self, inputstring: str, document: nodes.document) -> None:
         # open the notebook execution client,
         # this may execute the notebook immediately or during the page render
         with create_client(
-            notebook, document_path, nb_config, logger, nb_reader.read_fmt
+            notebook, document_path, nb_config, logger, nb_reader.read_fmt, km=self.km
         ) as nb_client:
             mdit_parser.options["nb_client"] = nb_client
             # convert to docutils AST, which is added to the document