Source code for myst_nb.core.execute.cache

"""Execute a notebook from the cache."""
from __future__ import annotations

from contextlib import nullcontext, suppress
from datetime import datetime
import os
from tempfile import TemporaryDirectory
from typing import ContextManager

from jupyter_cache import get_cache
from jupyter_cache.base import CacheBundleIn
from jupyter_cache.cache.db import NbProjectRecord
from jupyter_cache.executors.utils import single_nb_execution

from .base import ExecutionError, NotebookClientBase

[docs]class NotebookClientCache(NotebookClientBase): """A notebook client that retrieves notebook outputs from the cache, or executes the notebook and adds them to the cache, on entrance of the context. """ def start_client(self): # setup the cache cache = get_cache(self.nb_config.execution_cache_path or ".jupyter_cache") # TODO config on what notebook/cell metadata to hash/merge # attempt to match the notebook to one in the cache cache_record = None with suppress(KeyError): cache_record = cache.match_cache_notebook(self.notebook) # use the cached notebook if it exists if cache_record is not None:"Using cached notebook: ID={}") _, self._notebook = cache.merge_match_into_notebook(self.notebook) self.exec_metadata = { "mtime": cache_record.created.timestamp(), "runtime":"execution_seconds", None), "method": self.nb_config.execution_mode, "succeeded": True, "error": None, "traceback": None, } return if self.path is None: raise ValueError( "Input source must exist as file, if execution_mode is 'cache'" ) # attempt to execute the notebook read_fmt = self._kwargs.get("read_fmt", None) if read_fmt is not None: stage_record = cache.add_nb_to_project(str(self.path), read_data=read_fmt) else: stage_record = cache.add_nb_to_project(str(self.path)) # TODO do in try/except, in case of db write errors NbProjectRecord.remove_tracebacks([], cache.db) cwd_context: ContextManager[str] = ( TemporaryDirectory() # type: ignore if self.nb_config.execution_in_temp else nullcontext(str(self.path.parent)) ) with cwd_context as cwd: cwd = os.path.abspath(cwd) "Executing notebook using " + ("temporary" if self.nb_config.execution_in_temp else "local") + " CWD" ) result = single_nb_execution( self.notebook, cwd=cwd, allow_errors=self.nb_config.execution_allow_errors, timeout=self.nb_config.execution_timeout, meta_override=True, # TODO still support this? ) # handle success / failure cases # TODO do in try/except to be careful (in case of database write errors? if result.err is not None: msg = f"Executing notebook failed: {result.err.__class__.__name__}" if self.nb_config.execution_show_tb: msg += f"\n{result.exc_string}" self.logger.warning(msg, subtype="exec") if self.nb_config.execution_raise_on_error: raise ExecutionError(str(self.path)) from result.err NbProjectRecord.set_traceback(stage_record.uri, result.exc_string, cache.db) else:"Executed notebook in {result.time:.2f} seconds") cache_record = cache.cache_notebook_bundle( CacheBundleIn( self.notebook, stage_record.uri, data={"execution_seconds": result.time}, ), check_validity=False, overwrite=True, )"Cached executed notebook: ID={}") self.exec_metadata = { "mtime":, "runtime": result.time, "method": self.nb_config.execution_mode, "succeeded": False if result.err else True, "error": f"{result.err.__class__.__name__}" if result.err else None, "traceback": result.exc_string if result.err else None, }