From 4cb9d4a3ed1619da8d2996d68452c3831e69ddf6 Mon Sep 17 00:00:00 2001 From: Phil Schaf Date: Wed, 1 Jul 2026 09:18:08 +0200 Subject: [PATCH] feat: allow passing kernel manager to parser --- myst_nb/core/execute/__init__.py | 11 ++++++++--- myst_nb/core/execute/base.py | 8 +++++++- myst_nb/core/execute/cache.py | 1 + myst_nb/core/execute/direct.py | 1 + myst_nb/core/execute/inline.py | 1 + myst_nb/docutils_.py | 14 ++++++++++++-- myst_nb/sphinx_.py | 12 ++++++++++-- 7 files changed, 40 insertions(+), 8 deletions(-) diff --git a/myst_nb/core/execute/__init__.py b/myst_nb/core/execute/__init__.py index 5945988a..48deeed8 100644 --- a/myst_nb/core/execute/__init__.py +++ b/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) diff --git a/myst_nb/core/execute/base.py b/myst_nb/core/execute/base.py index 21570315..eaf6d5ba 100644 --- a/myst_nb/core/execute/base.py +++ b/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,6 +57,8 @@ def __init__( path: Path | None, nb_config: NbParserConfig, logger: LoggerType, + *, + km: KernelManager | None = None, **kwargs: Any, ): """Initialize the client.""" @@ -61,6 +66,7 @@ def __init__( self._path = path self._nb_config = nb_config self._logger = logger + self._km = km self._kwargs = kwargs self._glue_data: dict[str, NotebookNode] = {} diff --git a/myst_nb/core/execute/cache.py b/myst_nb/core/execute/cache.py index 9d5633e2..91e6f70c 100644 --- a/myst_nb/core/execute/cache.py +++ b/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 diff --git a/myst_nb/core/execute/direct.py b/myst_nb/core/execute/direct.py index e82624cb..13461f52 100644 --- a/myst_nb/core/execute/direct.py +++ b/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: diff --git a/myst_nb/core/execute/inline.py b/myst_nb/core/execute/inline.py index 06078d04..f0cedf6e 100644 --- a/myst_nb/core/execute/inline.py +++ b/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: diff --git a/myst_nb/docutils_.py b/myst_nb/docutils_.py index 06f80968..0ab051e7 100644 --- a/myst_nb/docutils_.py +++ b/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) diff --git a/myst_nb/sphinx_.py b/myst_nb/sphinx_.py index 14cec577..573e68c6 100644 --- a/myst_nb/sphinx_.py +++ b/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