from contextlib import contextmanager import copy import hashlib import io from pathlib import Path import shutil from typing import Iterable, List, Mapping, Optional, Tuple, Union import nbformat as nbf from jupyter_cache.base import ( # noqa: F401 NB_VERSION, CacheBundleIn, CacheBundleOut, CachingError, JupyterCacheAbstract, NbArtifactsAbstract, NbValidityError, ProjectNb, RetrievalError, ) from jupyter_cache.readers import DEFAULT_READ_DATA, NbReadError, get_reader from jupyter_cache.utils import to_relative_paths from .db import NbCacheRecord, NbProjectRecord, Setting, create_db, get_version CACHE_LIMIT_KEY = "cache_limit" DEFAULT_CACHE_LIMIT = 1000 class NbArtifacts(NbArtifactsAbstract): """Container for artefacts of a notebook execution.""" def __init__(self, paths: List[str], in_folder, check_existence=True): """Initiate NbArtifacts :param paths: list of paths :param check_existence: check the paths exist :param in_folder: The folder that all paths should be in (or subfolder). :raises IOError: if check_existence and file does not exist """ self.paths = [Path(p).absolute() for p in paths] self.in_folder = Path(in_folder).absolute() to_relative_paths(self.paths, self.in_folder, check_existence=check_existence) @property def relative_paths(self) -> List[Path]: """Return the list of paths (relative to the notebook folder).""" return to_relative_paths(self.paths, self.in_folder) def __iter__(self) -> Iterable[Tuple[Path, io.BufferedReader]]: """Yield the relative path and open files (in bytes mode)""" for path in self.paths: with path.open("rb") as handle: yield path.relative_to(self.in_folder), handle class JupyterCacheBase(JupyterCacheAbstract): def __init__(self, path): self._path = Path(path).absolute() self._db = None @property def path(self): if not self._path.exists(): self._path.mkdir(parents=True) return self._path @property def db(self): """a simple database for storing persistent global data.""" if self._db is None: self._db = create_db(self.path) return self._db def __repr__(self): return f"{self.__class__.__name__}({repr(str(self._path))})" def __getstate__(self): """For pickling instances, db must be removed.""" state = self.__dict__.copy() state["_db"] = None return state def get_version(self) -> Optional[str]: return get_version(self.path) def clear_cache(self): """Clear the cache completely.""" shutil.rmtree(self.path) self._db = None def _get_notebook_path_cache(self, hashkey, raise_on_missing=False) -> Path: """Retrieve a relative path in the cache to a notebook, from its hash.""" path = self.path.joinpath(Path("executed", hashkey, "base.ipynb")) if not path.exists() and raise_on_missing: raise RetrievalError(f"hashkey not in cache: {hashkey}") return path def _get_artifact_path_cache(self, hashkey) -> Path: """Retrieve a relative path in the cache to a notebook, from its hash.""" path = self.path.joinpath(Path("executed", hashkey, "artifacts")) return path def truncate_caches(self): """If the number of cached notebooks exceeds set limit, delete the oldest.""" cache_limit = Setting.get_value(CACHE_LIMIT_KEY, self.db, DEFAULT_CACHE_LIMIT) # TODO you could have better control over this by e.g. tagging certain caches # that should not be deleted. pks = NbCacheRecord.records_to_delete(cache_limit, self.db) for pk in pks: self.remove_cache(pk) def get_cache_limit(self): return Setting.get_value(CACHE_LIMIT_KEY, self.db, DEFAULT_CACHE_LIMIT) def change_cache_limit(self, size: int): assert isinstance(size, int) and size > 0 Setting.set_value(CACHE_LIMIT_KEY, size, self.db) def create_hashed_notebook( self, nb: nbf.NotebookNode, nb_metadata: Optional[Iterable[str]] = ("kernelspec",), cell_metadata: Optional[Iterable[str]] = None, ) -> Tuple[nbf.NotebookNode, str]: """Convert a notebook to a standard format and hash. Note: we always hash notebooks as version 4.4, to allow for matching notebooks of different versions :param nb_metadata: The notebook metadata keys to hash (if None, use all) :param cell_metadata: The cell metadata keys to hash (if None, use all) :return: (notebook, hash) """ # copy the notebook nb = copy.deepcopy(nb) # update the notebook to consistent version 4.4 nb = nbf.convert(nb, to_version=NB_VERSION) if nb.nbformat_minor > 5: raise CachingError("notebook version greater than 4.5 not yet supported") # remove non-code cells nb.cells = [cell for cell in nb.cells if cell.cell_type == "code"] # create notebook for hashing, with selected metadata hash_nb = nbf.from_dict( { "nbformat": nb.nbformat, "nbformat_minor": 4, # v4.5 include cell ids, which we do not cache "metadata": { k: v for k, v in nb.metadata.items() if nb_metadata is None or (k in nb_metadata) }, "cells": [ { "cell_type": cell.cell_type, "source": cell.source, "metadata": { k: v for k, v in cell.metadata.items() if cell_metadata is None or (k in cell_metadata) }, "execution_count": None, "outputs": [], } for cell in nb.cells if cell.cell_type == "code" ], } ) # hash notebook string = nbf.writes(hash_nb, nbf.NO_CONVERT) hash_string = hashlib.md5(string.encode()).hexdigest() return (nb, hash_string) def _validate_nb_bundle(self, nb_bundle: CacheBundleIn): """Validate that a notebook bundle should be cached. We check that the notebook has been executed correctly, by asserting `execution_count`s are consecutive and start at 1. """ execution_count = 1 for i, cell in enumerate(nb_bundle.nb.cells): if cell.cell_type != "code": continue if cell.execution_count != execution_count: raise NbValidityError( "Expected cell {} to have execution_count {} not {}".format( i, execution_count, cell.execution_count ), nb_bundle, ) execution_count += 1 # TODO check for output exceptions? # TODO assets def cache_notebook_bundle( self, bundle: CacheBundleIn, check_validity: bool = True, overwrite: bool = False, description="", ) -> NbCacheRecord: """Cache an executed notebook.""" # TODO it would be ideal to have some 'rollback' mechanism on exception if check_validity: self._validate_nb_bundle(bundle) hashed_nb, hashkey = self.create_hashed_notebook(bundle.nb) path = self._get_notebook_path_cache(hashkey) if path.exists(): if not overwrite: raise CachingError( "Notebook already exists in cache and overwrite=False." ) shutil.rmtree(path.parent) try: record = NbCacheRecord.record_from_hashkey(hashkey, self.db) except KeyError: pass else: NbCacheRecord.remove_record(record.pk, self.db) record = NbCacheRecord.create_record( uri=bundle.uri, hashkey=hashkey, db=self.db, data=bundle.data, description=description, ) path.parent.mkdir(parents=True) path.write_text(nbf.writes(hashed_nb, nbf.NO_CONVERT), encoding="utf8") # write artifacts artifact_folder = self._get_artifact_path_cache(hashkey) if artifact_folder.exists(): shutil.rmtree(artifact_folder) for rel_path, handle in bundle.artifacts or []: write_path = artifact_folder.joinpath(rel_path) write_path.parent.mkdir(parents=True, exist_ok=True) write_path.write_bytes(handle.read()) self.truncate_caches() return record def cache_notebook_file( self, path: str, uri: Optional[str] = None, artifacts: List[str] = (), data: Optional[dict] = None, check_validity: bool = True, overwrite: bool = False, ) -> NbCacheRecord: """Cache an executed notebook, returning its primary key. Note: non-code source text (e.g. markdown) is not stored in the cache. :param path: path to the notebook :param uri: alternative URI to store in the cache record (defaults to path) :param artifacts: list of paths to outputs of the executed notebook. Artifacts must be in the same folder as the notebook (or a sub-folder) :param data: additional, JSONable, data to store in the cache record :param check_validity: check that the notebook has been executed correctly, by asserting `execution_count`s are consecutive and start at 1. :param overwrite: Allow overwrite of cached notebooks with matching hash :return: The primary key of the cache record """ notebook = nbf.read(str(path), nbf.NO_CONVERT) return self.cache_notebook_bundle( CacheBundleIn( notebook, uri or str(path), artifacts=NbArtifacts(artifacts, in_folder=Path(path).parent), data=data or {}, ), check_validity=check_validity, overwrite=overwrite, ) def list_cache_records(self) -> List[NbCacheRecord]: return NbCacheRecord.records_all(self.db) def get_cache_record(self, pk: int) -> NbCacheRecord: return NbCacheRecord.record_from_pk(pk, self.db) def get_cache_bundle(self, pk: int) -> CacheBundleOut: record = NbCacheRecord.record_from_pk(pk, self.db) NbCacheRecord.touch(pk, self.db) path = self._get_notebook_path_cache(record.hashkey) artifact_folder = self._get_artifact_path_cache(record.hashkey) if not path.exists(): raise KeyError(f"Notebook file does not exist for cache record PK: {pk}") return CacheBundleOut( nbf.reads(path.read_text(encoding="utf8"), nbf.NO_CONVERT), record=record, artifacts=NbArtifacts( [p for p in artifact_folder.glob("**/*") if p.is_file()], in_folder=artifact_folder, ), ) @contextmanager def cache_artefacts_temppath(self, pk: int) -> Path: """Context manager to provide a temporary folder path to the notebook artifacts. Note this path is only guaranteed to exist within the scope of the context, and should only be used for read/copy operations:: with cache.cache_artefacts_temppath(1) as path: shutil.copytree(path, destination) """ record = NbCacheRecord.record_from_pk(pk, self.db) yield self._get_artifact_path_cache(record.hashkey) def remove_cache(self, pk: int): record = NbCacheRecord.record_from_pk(pk, self.db) path = self._get_notebook_path_cache(record.hashkey) if not path.exists(): raise KeyError(f"Notebook file does not exist for cache record PK: {pk}") shutil.rmtree(path.parent) NbCacheRecord.remove_records([pk], self.db) def match_cache_notebook(self, nb: nbf.NotebookNode) -> NbCacheRecord: """Match to an executed notebook, returning its primary key. :raises KeyError: if no match is found """ _, hashkey = self.create_hashed_notebook(nb) cache_record = NbCacheRecord.record_from_hashkey(hashkey, self.db) return cache_record def merge_match_into_notebook( self, nb: nbf.NotebookNode, nb_meta: Optional[Iterable[str]] = ("kernelspec", "language_info", "widgets"), cell_meta: Optional[Iterable[str]] = None, ) -> Tuple[int, nbf.NotebookNode]: """Match to an executed notebook and return a merged version :param nb: The input notebook :param nb_meta: metadata keys to merge from the cached notebook (all if None) :param cell_meta: cell metadata keys to merge from cached notebook (all if None) :raises KeyError: if no match is found :return: pk, input notebook with cached code cells and metadata merged. """ pk = self.match_cache_notebook(nb).pk cache_nb = self.get_cache_bundle(pk).nb nb = nbf.convert(copy.deepcopy(nb), NB_VERSION) if nb_meta is None: nb.metadata = cache_nb.metadata else: for key in nb_meta: if key in cache_nb.metadata: nb.metadata[key] = cache_nb.metadata[key] for idx in range(len(nb.cells)): if nb.cells[idx].cell_type == "code": cache_cell = cache_nb.cells.pop(0) in_cell = nb.cells[idx] if cell_meta is not None: # update the input metadata with select cached notebook metadata # then add the input metadata to the cached cell in_cell.metadata.update( {k: v for k, v in cache_cell.metadata.items() if k in cell_meta} ) cache_cell.metadata = in_cell.metadata if nb.nbformat_minor >= 5: cache_cell.id = in_cell.id else: cache_cell.pop("id", None) nb.cells[idx] = cache_cell return pk, nb def diff_nbnode_with_cache( self, pk: int, nb: nbf.NotebookNode, uri: str = "", as_str=False, **kwargs ): """Return a diff of a notebook to a cached one. Note: this will not diff markdown content, since it is not stored in the cache. """ try: import nbdime except ImportError: raise ImportError( "nbdime is required to diff notebooks, install with `pip install nbdime`" ) from nbdime.prettyprint import PrettyPrintConfig, pretty_print_diff cached_nb = self.get_cache_bundle(pk).nb nb, _ = self.create_hashed_notebook(nb) diff = nbdime.diff_notebooks(cached_nb, nb) if not as_str: return diff stream = io.StringIO() stream.writelines(["nbdiff\n", f"--- cached pk={pk}\n", f"+++ other: {uri}\n"]) pretty_print_diff( cached_nb, diff, "nb", PrettyPrintConfig(out=stream, **kwargs) ) return stream.getvalue() def add_nb_to_project( self, path: str, *, read_data: Mapping = DEFAULT_READ_DATA, assets: List[str] = (), ) -> NbProjectRecord: # check the reader can be loaded read_data = dict(read_data) _ = get_reader(read_data) # TODO should we test that the file can be read by the reader? return NbProjectRecord.create_record( str(Path(path).absolute()), self.db, raise_on_exists=False, read_data=read_data, assets=assets, ) # TODO physically copy to cache? # TODO assets def list_project_records( self, filter_uris: Optional[List[str]] = None, filter_pks: Optional[List[int]] = None, ) -> List[NbProjectRecord]: records = NbProjectRecord.records_all(self.db) if filter_uris is not None: records = [r for r in records if r.uri in filter_uris] if filter_pks is not None: records = [r for r in records if r.pk in filter_pks] return records def get_project_record(self, uri_or_pk: Union[int, str]) -> NbProjectRecord: if isinstance(uri_or_pk, int): record = NbProjectRecord.record_from_pk(uri_or_pk, self.db) else: record = NbProjectRecord.record_from_uri(uri_or_pk, self.db) return record def remove_nb_from_project(self, uri_or_pk: Union[int, str]): if isinstance(uri_or_pk, int): NbProjectRecord.remove_pks([uri_or_pk], self.db) else: NbProjectRecord.remove_uris([uri_or_pk], self.db) # TODO add discard all/multiple project records method def get_project_notebook(self, uri_or_pk: Union[int, str]) -> ProjectNb: if isinstance(uri_or_pk, int): record = NbProjectRecord.record_from_pk(uri_or_pk, self.db) else: record = NbProjectRecord.record_from_uri(uri_or_pk, self.db) if not Path(record.uri).exists(): raise OSError( f"The URI of the project record no longer exists: {record.uri}" ) try: reader = get_reader(record.read_data) notebook = reader(record.uri) assert isinstance( notebook, nbf.NotebookNode ), f"Reader did not return a v4 NotebookNode: {type(notebook)} {notebook}" except Exception as exc: raise NbReadError(f"Failed to read the notebook: {exc}") from exc return ProjectNb(record.pk, record.uri, notebook, record.assets) def get_cached_project_nb( self, uri_or_pk: Union[int, str] ) -> Optional[NbCacheRecord]: nb = self.get_project_notebook(uri_or_pk).nb _, hashkey = self.create_hashed_notebook(nb) try: return NbCacheRecord.record_from_hashkey(hashkey, self.db) except KeyError: return None def list_unexecuted( self, filter_uris: Optional[List[str]] = None, filter_pks: Optional[List[int]] = None, ) -> List[NbProjectRecord]: records = [] for record in self.list_project_records(filter_uris, filter_pks): nb = self.get_project_notebook(record.uri).nb _, hashkey = self.create_hashed_notebook(nb) try: NbCacheRecord.record_from_hashkey(hashkey, self.db) except KeyError: records.append(record) return records # removed until defined use case # def get_cache_codecell(self, pk: int, index: int) -> nbf.NotebookNode: # """Return a code cell from a cached notebook. # NOTE: the index **only** refers to the list of code cells, e.g. # `[codecell_0, textcell_1, codecell_2]` # would map {0: codecell_0, 1: codecell_2} # """ # nb_bundle = self.get_cache_bundle(pk) # _code_index = 0 # for cell in nb_bundle.nb.cells: # if cell.cell_type != "code": # continue # if _code_index == index: # return cell # _code_index += 1 # raise RetrievalError(f"Notebook contains less than {index+1} code cell(s)")