Source code for myst_nb.core.execute.base

"""Module for executing notebooks."""
from __future__ import annotations

from pathlib import Path
import re
from typing import Any

from nbformat import NotebookNode
from typing_extensions import TypedDict, final

from myst_nb.core.config import NbParserConfig
from myst_nb.core.loggers import LoggerType
from myst_nb.core.nb_to_tokens import nb_node_to_dict
from myst_nb.ext.glue import extract_glue_data

[docs]class ExecutionResult(TypedDict): """Result of executing a notebook.""" mtime: float """POSIX timestamp of the execution time""" runtime: float | None """runtime in seconds""" method: str """method used to execute the notebook""" succeeded: bool """True if the notebook executed successfully""" error: str | None """error type if the notebook failed to execute""" traceback: str | None """traceback if the notebook failed"""
[docs]class ExecutionError(Exception): """An exception for failed execution and `execution_raise_on_error` is true."""
[docs]class EvalNameError(Exception): """An exception for if an evaluation variable name is invalid."""
EVAL_NAME_REGEX = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
[docs]class NotebookClientBase: """A base client for interacting with Jupyter notebooks. This class is intended to be used as a context manager, and should only be entered once. Subclasses should override the `start_client` and `close_client` methods. """ def __init__( self, notebook: NotebookNode, path: Path | None, nb_config: NbParserConfig, logger: LoggerType, **kwargs: Any, ): """Initialize the client.""" self._notebook = notebook self._path = path self._nb_config = nb_config self._logger = logger self._kwargs = kwargs self._glue_data: dict[str, NotebookNode] = {} self._exec_metadata: ExecutionResult | None = None # get or create source map of cell to source line # use 1-based indexing rather than 0, or pseudo base of the cell index _source_map: list[int] = notebook.metadata.get("source_map", None) self._source_map = [ (_source_map[i] if _source_map else ((i + 1) * 10000)) + 1 for i, _ in enumerate(notebook.cells) ] @final def __enter__(self) -> NotebookClientBase: """Enter the context manager.""" self.start_client() # extract glue data from the notebook self._glue_data = extract_glue_data( self.notebook, self._source_map, self.logger ) return self @final def __exit__(self, exc_type, exc_val, exc_tb): """Exit the context manager.""" self.close_client(exc_type, exc_val, exc_tb)
[docs] def start_client(self): """Start the client."""
[docs] def finalise_client(self): """Finalise the client. This is called before final rendering and should be used, for example, to finalise the widget state on the metadata. """
[docs] def close_client(self, exc_type, exc_val, exc_tb): """Close the client."""
@property def notebook(self) -> NotebookNode: """Get the notebook.""" return self._notebook @property def path(self) -> Path | None: """Get the notebook path.""" return self._path @property def nb_config(self) -> NbParserConfig: """Get the notebook configuration.""" return self._nb_config @property def logger(self) -> LoggerType: """Get the logger.""" return self._logger @property def glue_data(self) -> dict[str, NotebookNode]: """Get the glue data.""" return self._glue_data @property def exec_metadata(self) -> ExecutionResult | None: """Get the execution metadata.""" return self._exec_metadata @exec_metadata.setter def exec_metadata(self, value: ExecutionResult): """Set the execution metadata.""" self._exec_metadata = value
[docs] def cell_line(self, cell_index: int) -> int: """Get the source line number of a cell.""" return self._source_map[cell_index]
@property def nb_metadata(self) -> dict[str, Any]: """Get the notebook level metadata.""" return nb_node_to_dict(self.notebook.get("metadata", {}))
[docs] def nb_source_code_lexer(self) -> str | None: """Get the lexer for the notebook source code.""" metadata = self.notebook.get("metadata", {}) langinfo = metadata.get("language_info") or {} lexer = langinfo.get("pygments_lexer") or langinfo.get("name", None) if lexer is None: lexer = (metadata.get("kernelspec") or {}).get("language", None) return lexer
[docs] def code_cell_outputs( self, cell_index: int ) -> tuple[int | None, list[NotebookNode]]: """Get the outputs of a cell. :returns: a tuple of the execution_count and the outputs :raises IndexError: if the cell index is out of range """ cells = self.notebook.get("cells", []) cell = cells[cell_index] return cell.get("execution_count", None), cell.get("outputs", [])
[docs] def eval_variable(self, name: str) -> NotebookNode | None: """Retrieve the value of a variable from the kernel. :param name: the name of the variable, must match the regex `[a-zA-Z][a-zA-Z0-9_]*` :returns: a code cell output :raises NotImplementedError: if the execution mode does not support this feature :raises EvalNameError: if the variable name is invalid """ raise NotImplementedError