Skip to content

Commit

Permalink
✨ NEW: Add "inline" execution mode
Browse files Browse the repository at this point in the history
chrisjsewell committed May 3, 2022
1 parent 8bd876c commit 86be2ab
Showing 14 changed files with 290 additions and 38 deletions.
1 change: 1 addition & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
@@ -50,6 +50,7 @@ repos:
- importlib_metadata
- myst-parser~=0.17.2
- "sphinx~=4.3.2"
- nbclient
- types-PyYAML
files: >
(?x)^(
4 changes: 2 additions & 2 deletions docs/index.md
Original file line number Diff line number Diff line change
@@ -72,8 +72,8 @@ Cache execution outputs, for fast re-builds.
:link: render/code-cells
:link-type: ref

Convert Jupyter execution outputs to embedded content.\
Insert outputs as variables into your documents.\
Convert Jupyter execution outputs to rich embedded content.\
Insert computed variables within your documents.\
Build single or collections of documents into multiple formats, including HTML websites and PDF books.

+++
4 changes: 2 additions & 2 deletions docs/reference/api.rst
Original file line number Diff line number Diff line change
@@ -49,10 +49,10 @@ Execute
:members:

.. autoclass:: myst_nb.core.execute.direct.NotebookClientDirect
:members:

.. autoclass:: myst_nb.core.execute.cache.NotebookClientCache
:members:

.. autoclass:: myst_nb.core.execute.inline.NotebookClientInline

.. autoclass:: myst_nb.core.execute.base.ExecutionResult
:members:
3 changes: 2 additions & 1 deletion myst_nb/core/config.py
Original file line number Diff line number Diff line change
@@ -187,7 +187,7 @@ def __post_init__(self):
"sections": (Section.global_lvl, Section.execute),
},
)
execution_mode: Literal["off", "force", "auto", "cache"] = dc.field(
execution_mode: Literal["off", "force", "auto", "cache", "inline"] = dc.field(
default="auto",
metadata={
"validator": in_(
@@ -196,6 +196,7 @@ def __post_init__(self):
"auto",
"force",
"cache",
"inline",
]
),
"help": "Execution mode for notebooks",
4 changes: 4 additions & 0 deletions myst_nb/core/execute/__init__.py
Original file line number Diff line number Diff line change
@@ -6,6 +6,7 @@
from .base import ExecutionError, ExecutionResult, NotebookClientBase # noqa: F401
from .cache import NotebookClientCache
from .direct import NotebookClientDirect
from .inline import NotebookClientInline

if TYPE_CHECKING:
from nbformat import NotebookNode
@@ -63,4 +64,7 @@ def create_client(
if nb_config.execution_mode == "cache":
return NotebookClientCache(notebook, path, nb_config, logger, read_fmt=read_fmt)

if nb_config.execution_mode == "inline":
return NotebookClientInline(notebook, path, nb_config, logger)

return NotebookClientBase(notebook, path, nb_config, logger)
3 changes: 1 addition & 2 deletions myst_nb/core/execute/base.py
Original file line number Diff line number Diff line change
@@ -58,7 +58,7 @@ def __init__(
self._logger = logger
self._kwargs = kwargs

self._glue_data: dict[str, NotebookNode] | None = None
self._glue_data: dict[str, NotebookNode] = {}
self._exec_metadata: ExecutionResult | None = None

# get or create source map of cell to source line
@@ -113,7 +113,6 @@ def logger(self) -> LoggerType:
@property
def glue_data(self) -> dict[str, NotebookNode]:
"""Get the glue data."""
assert self._glue_data is not None
return self._glue_data

@property
141 changes: 141 additions & 0 deletions myst_nb/core/execute/inline.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,141 @@
"""Execute a notebook inline."""
from __future__ import annotations

from datetime import datetime
import shutil
from tempfile import mkdtemp
import time
import traceback

from nbclient.client import CellExecutionError, CellTimeoutError, NotebookClient
from nbformat import NotebookNode

from myst_nb.glue import extract_glue_data_cell

from .base import ExecutionError, NotebookClientBase


class NotebookClientInline(NotebookClientBase):
"""A notebook client that executes the notebook inline,
i.e. during the render.
This allows for the client to be called in-between code cell executions,
in order to extract variable state.
"""

def start_client(self):

self._tmp_path = None
if self.nb_config.execution_in_temp:
self._tmp_path = mkdtemp()
resources = {"metadata": {"path": self._tmp_path}}
else:
if self.path is None:
raise ValueError(
"Input source must exist as file, if execution_in_temp=False"
)
resources = {"metadata": {"path": str(self.path.parent)}}

self.logger.info("Starting inline execution client")

self._time_start = time.perf_counter()
self._client = NotebookClient(
self.notebook,
record_timing=False,
resources=resources,
allow_errors=self.nb_config.execution_allow_errors,
timeout=self.nb_config.execution_timeout,
)
self._client.reset_execution_trackers()
if self._client.km is None:
self._client.km = self._client.create_kernel_manager()
if not self._client.km.has_kernel:
self._client.start_new_kernel()
self._client.start_new_kernel_client()

# retrieve the the language_info from the kernel
assert self._client.kc is not None
msg_id = self._client.kc.kernel_info()
info_msg = self._client.wait_for_reply(msg_id)
if info_msg is not None and "language_info" in info_msg["content"]:
self.notebook.metadata["language_info"] = info_msg["content"][
"language_info"
]
else:
self.logger.warning("Failed to retrieve language info from kernel")

self._last_cell_executed: int = -1
self._cell_error: None | Exception = None
self._exc_string: None | str = None

def close_client(self, exc_type, exc_val, exc_tb):
self.logger.info("Stopping inline execution client")
try:
# TODO because we set the widget state at the end,
# it won't be output by the renderer at present
self._client.set_widgets_metadata()
except Exception:
pass
if self._client.owns_km:
self._client._cleanup_kernel()
del self._client

_exec_time = time.perf_counter() - self._time_start
self.exec_metadata = {
"mtime": datetime.now().timestamp(),
"runtime": _exec_time,
"method": self.nb_config.execution_mode,
"succeeded": False if self._cell_error else True,
"error": f"{self._cell_error.__class__.__name__}"
if self._cell_error
else None,
"traceback": self._exc_string,
}
if not self._cell_error:
self.logger.info(f"Executed notebook in {_exec_time:.2f} seconds")
else:
msg = f"Executing notebook failed: {self._cell_error.__class__.__name__}"
if self.nb_config.execution_show_tb:
msg += f"\n{self._exc_string}"
self.logger.warning(msg, subtype="exec")
if self._tmp_path:
shutil.rmtree(self._tmp_path, ignore_errors=True)

def code_cell_outputs(
self, cell_index: int
) -> tuple[int | None, list[NotebookNode]]:

cells = self.notebook.get("cells", [])

# ensure all cells up to and including the requested cell have been executed
while (not self._cell_error) and cell_index > self._last_cell_executed:

self._last_cell_executed += 1
try:
next_cell = cells[self._last_cell_executed]
except IndexError:
break

try:
self._client.execute_cell(
next_cell,
self._last_cell_executed,
execution_count=self._client.code_cells_executed + 1,
)
except (CellExecutionError, CellTimeoutError) as err:
if self.nb_config.execution_raise_on_error:
raise ExecutionError(str(self.path)) from err
self._cell_error = err
self._exc_string = "".join(traceback.format_exc())

for key, cell_data in extract_glue_data_cell(next_cell):
if key in self._glue_data:
self.logger.warning(
f"glue key {key!r} duplicate",
subtype="glue",
line=self.cell_line(self._last_cell_executed),
)
self._glue_data[key] = cell_data

cell = cells[cell_index]
return cell.get("execution_count", None), cell.get("outputs", [])
52 changes: 31 additions & 21 deletions myst_nb/glue/__init__.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
"""Functionality for storing special data in notebook code cells,
which can then be inserted into the document body.
"""
from typing import Any, Dict, List
from __future__ import annotations

from typing import Any

import IPython
from IPython.display import display as ipy_display
@@ -12,7 +14,7 @@
GLUE_PREFIX = "application/papermill.record/"


def get_glue_roles(prefix: str = "glue:") -> Dict[str, Any]:
def get_glue_roles(prefix: str = "glue:") -> dict[str, Any]:
"""Return mapping of role names to role functions."""
from .roles import PasteMarkdownRole, PasteRoleAny, PasteTextRole

@@ -24,7 +26,7 @@ def get_glue_roles(prefix: str = "glue:") -> Dict[str, Any]:
}


def get_glue_directives(prefix: str = "glue:") -> Dict[str, Any]:
def get_glue_directives(prefix: str = "glue:") -> dict[str, Any]:
"""Return mapping of directive names to directive functions."""
from .directives import (
PasteAnyDirective,
@@ -42,7 +44,7 @@ def get_glue_directives(prefix: str = "glue:") -> Dict[str, Any]:
}


def glue(name: str, variable, display: bool = True) -> None:
def glue(name: str, variable: Any, display: bool = True) -> None:
"""Glue a variable into the notebook's cell metadata.
Parameters
@@ -68,34 +70,42 @@ def glue(name: str, variable, display: bool = True) -> None:

def extract_glue_data(
notebook: NotebookNode,
source_map: List[int],
source_map: list[int],
logger: LoggerType,
) -> Dict[str, NotebookNode]:
) -> dict[str, NotebookNode]:
"""Extract all the glue data from the notebook."""
# note this assumes v4 notebook format
data: Dict[str, NotebookNode] = {}
data: dict[str, NotebookNode] = {}
for index, cell in enumerate(notebook.cells):
if cell.cell_type != "code":
continue
outputs = []
for output in cell.get("outputs", []):
meta = output.get("metadata", {})
if "scrapbook" not in meta:
outputs.append(output)
continue
key = meta["scrapbook"]["name"]
mime_prefix = len(meta["scrapbook"].get("mime_prefix", ""))
for key, cell_data in extract_glue_data_cell(cell):
if key in data:
logger.warning(
f"glue key {key!r} duplicate",
subtype="glue",
line=source_map[index],
)
output["data"] = {k[mime_prefix:]: v for k, v in output["data"].items()}
data[key] = output
if not mime_prefix:
# assume that the output is a displayable object
outputs.append(output)
cell.outputs = outputs
data[key] = cell_data

return data


def extract_glue_data_cell(cell: NotebookNode) -> list[tuple[str, NotebookNode]]:
"""Extract glue data from a single cell."""
outputs = []
data = []
for output in cell.get("outputs", []):
meta = output.get("metadata", {})
if "scrapbook" not in meta:
outputs.append(output)
continue
key = meta["scrapbook"]["name"]
mime_prefix = len(meta["scrapbook"].get("mime_prefix", ""))
output["data"] = {k[mime_prefix:]: v for k, v in output["data"].items()}
data.append((key, output))
if not mime_prefix:
# assume that the output is a displayable object
outputs.append(output)
cell.outputs = outputs
return data
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -38,6 +38,7 @@ dependencies = [
"importlib_metadata",
"ipython",
"jupyter-cache~=0.5.0",
"nbclient", # nbclient version pinned by jupyter-client
"myst-parser~=0.17.2",
"nbformat~=5.0",
"pyyaml",
39 changes: 30 additions & 9 deletions tests/test_execute.py
Original file line number Diff line number Diff line change
@@ -63,6 +63,21 @@ def test_basic_unrun_cache(sphinx_run, file_regression, check_nbs):
assert data["succeeded"] is True


@pytest.mark.sphinx_params("basic_unrun.ipynb", conf={"nb_execution_mode": "inline"})
def test_basic_unrun_inline(sphinx_run, file_regression, check_nbs):
"""The outputs should be populated."""
sphinx_run.build()
assert sphinx_run.warnings() == ""
assert "test_name" in sphinx_run.app.env.metadata["basic_unrun"]
regress_nb_doc(file_regression, sphinx_run, check_nbs)

assert NbMetadataCollector.new_exec_data(sphinx_run.env)
data = NbMetadataCollector.get_exec_data(sphinx_run.env, "basic_unrun")
assert data
assert data["method"] == "inline"
assert data["succeeded"] is True


@pytest.mark.sphinx_params("basic_unrun.ipynb", conf={"nb_execution_mode": "cache"})
def test_rebuild_cache(sphinx_run):
"""The notebook should only be executed once."""
@@ -130,6 +145,21 @@ def test_basic_failing_auto(sphinx_run, file_regression, check_nbs):
sphinx_run.get_report_file()


@pytest.mark.skipif(ipy_version[0] < 8, reason="Error message changes for ipython v8")
@pytest.mark.sphinx_params("basic_failing.ipynb", conf={"nb_execution_mode": "inline"})
def test_basic_failing_inline(sphinx_run, file_regression, check_nbs):
sphinx_run.build()
assert "Executing notebook failed" in sphinx_run.warnings()
regress_nb_doc(file_regression, sphinx_run, check_nbs)

data = NbMetadataCollector.get_exec_data(sphinx_run.env, "basic_failing")
assert data
assert data["method"] == "inline"
assert data["succeeded"] is False
assert data["traceback"]
sphinx_run.get_report_file()


@pytest.mark.skipif(ipy_version[0] < 8, reason="Error message changes for ipython v8")
@pytest.mark.sphinx_params(
"basic_failing.ipynb",
@@ -172,15 +202,6 @@ def test_raise_on_error_cache(sphinx_run):
sphinx_run.build()


@pytest.mark.sphinx_params("basic_unrun.ipynb", conf={"nb_execution_mode": "force"})
def test_outputs_present(sphinx_run, file_regression, check_nbs):
sphinx_run.build()
# print(sphinx_run.status())
assert sphinx_run.warnings() == ""
assert "test_name" in sphinx_run.app.env.metadata["basic_unrun"]
regress_nb_doc(file_regression, sphinx_run, check_nbs)


@pytest.mark.sphinx_params(
"complex_outputs_unrun.ipynb", conf={"nb_execution_mode": "cache"}
)
Loading
Oops, something went wrong.

0 comments on commit 86be2ab

Please sign in to comment.