Coverage for databooks/data_models/notebook.py: 92%
113 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-20 15:04 +0200
1"""Data models - Jupyter Notebooks and components."""
3from __future__ import annotations
5import json
6from copy import deepcopy
7from difflib import SequenceMatcher
8from itertools import chain
9from pathlib import Path
10from typing import (
11 Any,
12 Callable,
13 Generator,
14 Iterable,
15 List,
16 Optional,
17 Sequence,
18 Tuple,
19 TypeVar,
20 Union,
21 cast,
22)
24from pydantic import RootModel
25from rich import box
26from rich.columns import Columns
27from rich.console import Console, ConsoleOptions, Group, RenderableType, RenderResult
28from rich.panel import Panel
29from rich.text import Text
31from databooks.data_models.base import BaseCells, DatabooksBase
32from databooks.data_models.cell import CellMetadata, CodeCell, MarkdownCell, RawCell
33from databooks.logging import get_logger
35logger = get_logger(__file__)
37Cell = Union[CodeCell, RawCell, MarkdownCell]
38CellsPair = Tuple[List[Cell], List[Cell]]
39T = TypeVar("T", Cell, CellsPair)
42class Cells(RootModel[Sequence[T]], BaseCells[T]):
43 """Similar to `list`, with `-` operator using `difflib.SequenceMatcher`."""
45 root: Sequence[T]
47 # `UserList.data` is a plain attribute; here it is derived from `root`
48 @property
49 def data(self) -> List[T]: # ty: ignore[invalid-property-type-override]
50 """Define property `data` required for `collections.UserList` class."""
51 return list(self.root)
53 def __iter__(self) -> Generator[Any, None, None]:
54 """Use list property as iterable."""
55 return (el for el in self.data)
57 def __sub__(self: Cells[Cell], other: Cells[Cell]) -> Cells[CellsPair]:
58 """Return the difference using `difflib.SequenceMatcher`."""
59 if type(self) is not type(other):
60 raise TypeError(
61 f"Unsupported operand types for `-`: `{type(self).__name__}` and"
62 f" `{type(other).__name__}`"
63 )
65 # By setting the context to the max number of cells and using
66 # `pathlib.SequenceMatcher.get_grouped_opcodes` we essentially get the same
67 # result as `pathlib.SequenceMatcher.get_opcodes` but in smaller chunks
68 n_context = max(len(self), len(other))
69 diff_opcodes = list(
70 SequenceMatcher(
71 isjunk=None, a=self, b=other, autojunk=False
72 ).get_grouped_opcodes(n_context)
73 )
75 if len(diff_opcodes) > 1:
76 raise RuntimeError(
77 "Expected one group for opcodes when context size is "
78 f" {n_context} for {len(self)} and {len(other)} cells in"
79 " notebooks."
80 )
82 # ty reads typeshed's `UserList.data` as a pydantic field of `Cells` and
83 # expects it in `__init__`, so constructing `Cells` needs this suppression
84 return Cells[CellsPair](
85 [
86 tuple((self.data[i1:j1], other.data[i2:j2]))
87 for _, i1, j1, i2, j2 in chain.from_iterable(diff_opcodes)
88 ]
89 ) # ty: ignore[missing-argument]
91 def __rich_console__(
92 self, console: Console, options: ConsoleOptions
93 ) -> RenderResult:
94 """Rich display of all cells in notebook."""
95 yield from self._get_renderables(expand=True, width=options.max_width // 3)
97 def _get_renderables(self, **wrap_cols_kwargs: Any) -> Iterable[RenderableType]:
98 """Get the Rich renderables, depending on whether `Cells` is a diff or not."""
99 if all(isinstance(el, tuple) for el in self.data):
100 return chain.from_iterable(
101 (
102 Cells.wrap_cols(val[0], val[1], **wrap_cols_kwargs)
103 if val[0] != val[1]
104 else val[0]
105 )
106 for val in cast(List[CellsPair], self.data)
107 )
108 return cast(List[Cell], self.data)
110 @classmethod
111 def __get_validators__(cls) -> Generator[Callable[..., Any], None, None]:
112 """Get validators for custom class."""
113 yield cls.validate
115 # pydantic v1 validator hook, narrower than `BaseModel.validate` on purpose
116 @classmethod
117 def validate(cls, v: List[T]) -> Cells[T]: # ty: ignore[invalid-method-override]
118 """Ensure object is custom defined container."""
119 if not isinstance(v, cls):
120 return cls(v) # ty: ignore[missing-argument]
121 else:
122 return v
124 @classmethod
125 def wrap_cols(
126 cls, first_cells: List[Cell], last_cells: List[Cell], **cols_kwargs: Any
127 ) -> Sequence[Columns]:
128 """Wrap the first and second cells into colunmns for iterable."""
129 _empty = [Panel(Text("<None>", justify="center"), box=box.SIMPLE)]
130 _first = Group(*first_cells or _empty)
131 _last = Group(*last_cells or _empty)
132 return [Columns([_first, _last], **cols_kwargs)]
134 @staticmethod
135 def wrap_git(
136 first_cells: List[Cell],
137 last_cells: List[Cell],
138 hash_first: Optional[str] = None,
139 hash_last: Optional[str] = None,
140 ) -> Sequence[Cell]:
141 """Wrap git-diff cells in existing notebook."""
142 return [
143 MarkdownCell(
144 metadata=CellMetadata(git_hash=hash_first),
145 source=[f"`<<<<<<< {hash_first}`"],
146 ),
147 *first_cells,
148 MarkdownCell(
149 source=["`=======`"],
150 metadata=CellMetadata(),
151 ),
152 *last_cells,
153 MarkdownCell(
154 metadata=CellMetadata(git_hash=hash_last),
155 source=[f"`>>>>>>> {hash_last}`"],
156 ),
157 ]
159 # Only diff cells (Cells[CellsPair]) can be resolved: narrower self than BaseCells
160 def resolve( # ty: ignore[invalid-method-override]
161 self: Cells[CellsPair],
162 *,
163 keep_first_cells: Optional[bool] = None,
164 first_id: Optional[str] = None,
165 last_id: Optional[str] = None,
166 **kwargs: Any,
167 ) -> List[Cell]:
168 """
169 Resolve differences between `databooks.data_models.notebook.Cells`.
171 :param keep_first_cells: Whether to keep the cells of the first notebook or not.
172 If `None`, then keep both wrapping the git-diff tags
173 :param first_id: Git hash of first file in conflict
174 :param last_id: Git hash of last file in conflict
175 :param kwargs: (Unused) keyword arguments to keep compatibility with
176 `databooks.data_models.base.resolve`
177 :return: List of cells
178 """
179 if keep_first_cells is not None:
180 return list(
181 chain.from_iterable(pairs[not keep_first_cells] for pairs in self.data)
182 )
183 return list(
184 chain.from_iterable(
185 (
186 Cells.wrap_git(
187 first_cells=val[0],
188 last_cells=val[1],
189 hash_first=first_id,
190 hash_last=last_id,
191 )
192 if val[0] != val[1]
193 else val[0]
194 )
195 for val in self.data
196 )
197 )
200class NotebookMetadata(DatabooksBase):
201 """Notebook metadata. Empty by default but can accept extra fields."""
204class JupyterNotebook(DatabooksBase, extra="forbid"):
205 """Jupyter notebook. Extra fields yield invalid notebook."""
207 nbformat: int
208 nbformat_minor: int
209 metadata: NotebookMetadata
210 cells: Cells[Cell]
212 def __rich_console__(
213 self, console: Console, options: ConsoleOptions
214 ) -> RenderResult:
215 """Rich display notebook."""
217 def _rich(kernel: str) -> Text:
218 """Display with `kernel` theme, horizontal padding and right-justified."""
219 return Text(kernel, style="kernel", justify="right")
221 kernelspec = self.metadata.model_dump().get("kernelspec", {})
222 if isinstance(kernelspec, tuple): # check if this is a `DiffCells`
223 kernelspec = tuple(
224 ks or {"language": "text", "display_name": "null"} for ks in kernelspec
225 )
226 lang_first, lang_last = (ks.get("language", "text") for ks in kernelspec)
227 nb_lang = lang_first if lang_first == lang_last else "text"
228 if any("display_name" in ks.keys() for ks in kernelspec):
229 kernel_first, kernel_last = [
230 _rich(ks["display_name"]) for ks in kernelspec
231 ]
232 yield (
233 Columns(
234 [kernel_first, kernel_last],
235 expand=True,
236 width=options.max_width // 3,
237 )
238 if kernel_first != kernel_last
239 else kernel_first
240 )
241 else:
242 nb_lang = kernelspec.get("language", "text")
243 if "display_name" in kernelspec.keys():
244 yield _rich(kernelspec["display_name"])
246 for cell in self.cells:
247 if isinstance(cell, CodeCell):
248 cell.metadata = CellMetadata(**cell.metadata.model_dump(), lang=nb_lang)
249 yield self.cells
251 @classmethod
252 def parse_file(cls, path: Path | str, **parse_kwargs: Any) -> JupyterNotebook:
253 """Parse notebook from a path."""
254 content_arg = parse_kwargs.pop("content_type", None)
255 if content_arg is not None:
256 raise ValueError(
257 f"Value of `content_type` must be `json` (default), got `{content_arg}`"
258 )
260 path = Path(path) if not isinstance(path, Path) else path
261 return JupyterNotebook.model_validate_json(
262 json_data=path.read_text(encoding="utf-8")
263 )
265 def write(
266 self, path: Path | str, overwrite: bool = False, **json_kwargs: Any
267 ) -> None:
268 """Write notebook to disk."""
269 path = Path(path) if not isinstance(path, Path) else path
270 json_kwargs = {"indent": 2, **json_kwargs}
271 if path.is_file() and not overwrite:
272 raise ValueError(
273 f"File exists at {path} exists. Specify `overwrite = True`."
274 )
276 self.__class__.model_validate(self.model_dump())
278 with path.open("w", encoding="utf-8") as f:
279 json.dump(self.model_dump(), fp=f, **json_kwargs)
281 def clear_metadata(
282 self,
283 *,
284 notebook_metadata_keep: Optional[Sequence[str]] = None,
285 notebook_metadata_remove: Optional[Sequence[str]] = None,
286 **cell_kwargs: Any,
287 ) -> None:
288 """
289 Clear notebook and cell metadata.
291 :param notebook_metadata_keep: Metadata values to keep - simply pass an empty
292 sequence (i.e.: `()`) to remove all extra fields.
293 :param notebook_metadata_remove: Metadata values to remove
294 :param cell_kwargs: keyword arguments to be passed to each cell's
295 `databooks.data_models.cell.BaseCell.clear_metadata`
296 :return:
297 """
298 nargs = sum(
299 (notebook_metadata_keep is not None, notebook_metadata_remove is not None)
300 )
301 if nargs != 1:
302 raise ValueError(
303 "Exactly one of `notebook_metadata_keep` or `notebook_metadata_remove`"
304 f" must be passed, got {nargs} arguments."
305 )
306 if notebook_metadata_keep is not None:
307 notebook_metadata_remove = tuple(
308 field
309 for field, _ in self.metadata
310 if field not in notebook_metadata_keep
311 )
312 # Non-None here: exactly one of keep/remove is given and keep became remove
313 self.metadata.remove_fields(cast(Sequence[str], notebook_metadata_remove))
315 if len(cell_kwargs) > 0:
316 _clean_cells = deepcopy(self.cells)
317 for cell in _clean_cells:
318 cell.clear_fields(**cell_kwargs)
319 self.cells = _clean_cells