Coverage for databooks/cli.py: 90%
124 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-20 15:04 +0200
1"""Main CLI application."""
3from itertools import compress
4from pathlib import Path
5from typing import List, Optional, Tuple
7import tomli
8from rich.progress import (
9 BarColumn,
10 Progress,
11 SpinnerColumn,
12 TextColumn,
13 TimeElapsedColumn,
14)
15from rich.prompt import Confirm
16from typer import Argument, BadParameter, Context, Exit, Option, Typer, echo
18from databooks.affirm import affirm_all
19from databooks.common import expand_paths
20from databooks.config import TOML_CONFIG_FILE, get_config
21from databooks.conflicts import conflicts2nbs, path2conflicts
22from databooks.git_utils import get_nb_diffs
23from databooks.logging import get_logger
24from databooks.metadata import clear_all
25from databooks.recipes import Recipe
26from databooks.tui import ImgFmt, diffs2rich, nbs2rich
27from databooks.version import __version__
29logger = get_logger(__file__)
31app = Typer()
34def _version_callback(show_version: bool) -> None:
35 """Return application version."""
36 if show_version:
37 echo("databooks version: " + __version__)
38 raise Exit()
41def _help_callback(ctx: Context, show_help: Optional[bool]) -> None:
42 """Reimplement `help` command to execute eagerly."""
43 if show_help:
44 echo(ctx.command.get_help(ctx))
45 raise Exit()
48def _config_callback(ctx: Context, config_path: Optional[Path]) -> Optional[Path]:
49 """Get config file and inject values into context to override default args."""
50 target_paths = expand_paths(
51 paths=[Path(p).resolve() for p in ctx.params.get("paths", ())]
52 ) or [Path.cwd()]
53 config_path = (
54 get_config(
55 target_paths=target_paths,
56 config_filename=TOML_CONFIG_FILE,
57 )
58 if config_path is None and target_paths
59 else config_path
60 )
61 logger.debug(f"Loading config file from: {config_path}")
62 if config_path is not None: # config may not be specified
63 with config_path.open("rb") as f:
64 conf = (
65 tomli.load(f)
66 .get("tool", {})
67 .get("databooks", {})
68 .get(ctx.command.name, {})
69 )
70 # Merge configuration
71 ctx.default_map = {
72 **(ctx.default_map or {}),
73 **{k.replace("-", "_"): v for k, v in conf.items()},
74 }
75 return config_path
78def _check_paths(paths: List[Path], ignore: List[str]) -> List[Path]:
79 """Check that notebooks exist retrieve the file paths."""
80 if any(path.suffix not in ("", ".ipynb") for path in paths):
81 raise BadParameter(
82 "Expected either notebook files, a directory or glob expression."
83 )
84 nb_paths = expand_paths(paths=paths, ignore=ignore)
85 if not nb_paths:
86 logger.info(
87 f"No notebooks found in {[p.resolve() for p in paths]}. Nothing to do."
88 )
89 raise Exit()
90 return nb_paths
93def _parse_paths(
94 *refs: Optional[str], paths: Optional[List[Path]]
95) -> Tuple[Tuple[Optional[str], ...], List[Path]]:
96 """Detect paths from `refs` and add to `paths`."""
97 paths = list(paths or ())
98 first, *rest = refs
99 if first is not None and Path(first).exists():
100 paths += [Path(first)]
101 first = None
102 if rest:
103 _refs, _paths = _parse_paths(*rest, paths=paths)
104 return (first, *_refs), _paths
105 return (first,), paths
108@app.callback()
109def callback( # noqa: D103
110 version: Optional[bool] = Option(
111 None, "--version", callback=_version_callback, is_eager=True
112 ),
113) -> None:
114 """CLI tool to resolve git conflicts and remove metadata in notebooks."""
117@app.command(add_help_option=False)
118def meta(
119 paths: List[Path] = Argument(..., is_eager=True, help="Path(s) of notebook files"),
120 ignore: List[str] = Option(["!*"], help="Glob expression(s) of files to ignore"),
121 prefix: str = Option("", help="Prefix to add to filepath when writing files"),
122 suffix: str = Option("", help="Suffix to add to filepath when writing files"),
123 rm_outs: bool = Option(False, help="Whether to remove cell outputs"),
124 rm_exec: bool = Option(True, help="Whether to remove the cell execution counts"),
125 nb_meta_keep: List[str] = Option((), help="Notebook metadata fields to keep"),
126 cell_meta_keep: List[str] = Option((), help="Cells metadata fields to keep"),
127 cell_fields_keep: List[str] = Option(
128 (),
129 help="Other (excluding `execution_counts` and `outputs`) cell fields to keep",
130 ),
131 overwrite: bool = Option(False, "--yes", "-y", help="Confirm overwrite of files"),
132 check: bool = Option(
133 False,
134 "--check",
135 help="Don't write files but check whether there is unwanted metadata",
136 ),
137 verbose: bool = Option(
138 False, "--verbose", "-v", help="Log processed files in console"
139 ),
140 config: Optional[Path] = Option(
141 None,
142 "--config",
143 "-c",
144 is_eager=True,
145 callback=_config_callback,
146 resolve_path=True,
147 exists=True,
148 help="Get CLI options from configuration file",
149 ),
150 help: Optional[bool] = Option(
151 None,
152 "--help",
153 is_eager=True,
154 callback=_help_callback,
155 help="Show this message and exit",
156 ),
157) -> None:
158 """Clear both notebook and cell metadata."""
159 nb_paths = _check_paths(paths=paths, ignore=ignore)
161 if not bool(prefix + suffix) and not check:
162 overwrite = (
163 Confirm.ask(
164 f"{len(nb_paths)} files will be overwritten"
165 " (no prefix nor suffix was passed). Continue?"
166 )
167 if not overwrite
168 else overwrite
169 )
170 if not overwrite:
171 raise Exit()
172 else:
173 logger.warning(f"{len(nb_paths)} files will be overwritten")
175 write_paths = [p.parent / (prefix + p.stem + suffix + p.suffix) for p in nb_paths]
176 cell_fields_keep = list(
177 compress(["outputs", "execution_count"], (not v for v in (rm_outs, rm_exec)))
178 ) + list(cell_fields_keep)
179 with Progress(
180 SpinnerColumn(),
181 TextColumn("[progress.description]{task.description}"),
182 BarColumn(),
183 TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
184 TimeElapsedColumn(),
185 ) as progress:
186 metadata = progress.add_task("[yellow]Removing metadata", total=len(nb_paths))
188 are_equal = clear_all(
189 read_paths=nb_paths,
190 write_paths=write_paths,
191 progress_callback=lambda: progress.update(metadata, advance=1),
192 notebook_metadata_keep=nb_meta_keep,
193 cell_metadata_keep=cell_meta_keep,
194 cell_fields_keep=cell_fields_keep,
195 check=check,
196 verbose=verbose,
197 overwrite=overwrite,
198 )
199 if check:
200 if all(are_equal):
201 logger.info("No unwanted metadata!")
202 else:
203 logger.info(
204 f"Found unwanted metadata in {sum(not eq for eq in are_equal)} out of"
205 f" {len(are_equal)} files."
206 )
207 raise Exit(code=1)
208 else:
209 logger.info(
210 f"The metadata of {sum(not eq for eq in are_equal)} out of {len(are_equal)}"
211 " notebooks were removed!"
212 )
215@app.command("assert", add_help_option=False)
216def affirm_meta(
217 paths: List[Path] = Argument(..., is_eager=True, help="Path(s) of notebook files"),
218 ignore: List[str] = Option(["!*"], help="Glob expression(s) of files to ignore"),
219 expr: List[str] = Option(
220 (), "--expr", "-x", help="Expressions to assert on notebooks"
221 ),
222 recipe: List[Recipe] = Option(
223 (),
224 "--recipe",
225 "-r",
226 help="Common recipes of expressions - see"
227 " https://databooks.dev/latest/usage/overview/#recipes",
228 ),
229 verbose: bool = Option(
230 False, "--verbose", "-v", help="Log processed files in console"
231 ),
232 config: Optional[Path] = Option(
233 None,
234 "--config",
235 "-c",
236 is_eager=True,
237 callback=_config_callback,
238 resolve_path=True,
239 exists=True,
240 help="Get CLI options from configuration file",
241 ),
242 help: Optional[bool] = Option(
243 None,
244 "--help",
245 is_eager=True,
246 callback=_help_callback,
247 help="Show this message and exit",
248 ),
249) -> None:
250 """
251 Assert notebook metadata has desired values.
253 Pass one (or multiple) strings or recipes. The available variables in scope include
254 `nb` (notebook), `raw_cells` (notebook cells of `raw` type), `md_cells` (notebook
255 cells of `markdown` type), `code_cells` (notebook cells of `code` type) and
256 `exec_cells` (notebook cells of `code` type that were executed - have an `execution
257 count` value). Recipes can be found on `databooks.recipes.CookBook`.
258 """
259 nb_paths = _check_paths(paths=paths, ignore=ignore)
260 exprs = [r.name for r in recipe] + list(expr)
261 if not exprs:
262 raise BadParameter("Must specify at least one of `expr` or `recipe`.")
264 with Progress(
265 SpinnerColumn(),
266 TextColumn("[progress.description]{task.description}"),
267 BarColumn(),
268 TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
269 TimeElapsedColumn(),
270 ) as progress:
271 assert_checks = progress.add_task(
272 "[yellow]Running assert checks", total=len(nb_paths)
273 )
275 are_ok = affirm_all(
276 nb_paths=nb_paths,
277 progress_callback=lambda: progress.update(assert_checks, advance=1),
278 exprs=exprs,
279 verbose=verbose,
280 )
282 if all(are_ok):
283 logger.info("All notebooks comply with the desired metadata!")
284 else:
285 logger.info(
286 f"Found issues in notebook metadata for {sum(not ok for ok in are_ok)} out"
287 f" of {len(are_ok)} notebooks."
288 )
289 raise Exit(code=1)
292@app.command(add_help_option=False)
293def fix(
294 paths: List[Path] = Argument(
295 ..., is_eager=True, help="Path(s) of notebook files with conflicts"
296 ),
297 ignore: List[str] = Option(["!*"], help="Glob expression(s) of files to ignore"),
298 metadata_head: bool = Option(
299 True, help="Whether or not to keep the metadata from the head/current notebook"
300 ),
301 cells_head: Optional[bool] = Option(
302 None,
303 help="Whether to keep the cells from the head/base notebook. Omit to keep both",
304 ),
305 cell_fields_ignore: List[str] = Option(
306 [
307 "id",
308 "execution_count",
309 ],
310 help="Cell fields to remove before comparing cells",
311 ),
312 interactive: bool = Option(
313 False,
314 "--interactive",
315 "-i",
316 help="Interactively resolve the conflicts (not implemented)",
317 ),
318 verbose: bool = Option(
319 False, "--verbose", "-v", help="Log processed files in console"
320 ),
321 config: Optional[Path] = Option(
322 None,
323 "--config",
324 "-c",
325 is_eager=True,
326 callback=_config_callback,
327 resolve_path=True,
328 exists=True,
329 help="Get CLI options from configuration file",
330 ),
331 help: Optional[bool] = Option(
332 None,
333 "--help",
334 is_eager=True,
335 callback=_help_callback,
336 help="Show this message and exit",
337 ),
338) -> None:
339 """
340 Fix git conflicts for notebooks.
342 Perform by getting the unmerged blobs from git index, comparing them and returning
343 a valid notebook summarizing the differences - see
344 [git docs](https://git-scm.com/docs/git-ls-files).
345 """
346 filepaths = expand_paths(paths=paths, ignore=ignore)
347 if filepaths is None:
348 raise RuntimeError(
349 f"Expected `filepaths` to be list of paths, got {filepaths}."
350 )
351 conflict_files = path2conflicts(nb_paths=filepaths)
352 if not conflict_files:
353 raise BadParameter(
354 f"No conflicts found at {', '.join([str(p) for p in filepaths])}."
355 )
356 if interactive:
357 raise NotImplementedError
359 with Progress(
360 SpinnerColumn(),
361 TextColumn("[progress.description]{task.description}"),
362 BarColumn(),
363 TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
364 TimeElapsedColumn(),
365 ) as progress:
366 conflicts = progress.add_task(
367 "[yellow]Resolving conflicts", total=len(conflict_files)
368 )
369 conflicts2nbs(
370 conflict_files=conflict_files,
371 meta_first=metadata_head,
372 cells_first=cells_head,
373 cell_fields_ignore=cell_fields_ignore,
374 verbose=verbose,
375 progress_callback=lambda: progress.update(conflicts, advance=1),
376 )
377 logger.info(f"Resolved the conflicts of {len(conflict_files)}!")
380@app.command(add_help_option=False)
381def show(
382 paths: List[Path] = Argument(
383 ..., is_eager=True, help="Path(s) of notebook files with conflicts"
384 ),
385 ignore: List[str] = Option(["!*"], help="Glob expression(s) of files to ignore"),
386 export: Optional[ImgFmt] = Option(
387 None,
388 "--export",
389 "-x",
390 help="Export rich outputs as a string.",
391 ),
392 pager: bool = Option(
393 False, "--pager", "-p", help="Use pager instead of printing to terminal"
394 ),
395 verbose: bool = Option(
396 False, "--verbose", "-v", help="Increase verbosity for debugging"
397 ),
398 multiple: bool = Option(False, "--yes", "-y", help="Show multiple files"),
399 config: Optional[Path] = Option(
400 None,
401 "--config",
402 "-c",
403 is_eager=True,
404 callback=_config_callback,
405 resolve_path=True,
406 exists=True,
407 help="Get CLI options from configuration file",
408 ),
409 help: Optional[bool] = Option(
410 None,
411 "--help",
412 is_eager=True,
413 callback=_help_callback,
414 help="Show this message and exit",
415 ),
416) -> None:
417 """Show rich representation of notebook."""
418 if export is not None and pager:
419 raise BadParameter("Cannot use both pager and export output.")
420 nb_paths = _check_paths(paths=paths, ignore=ignore)
421 if len(nb_paths) > 1 and not multiple:
422 if not Confirm.ask(f"Show {len(nb_paths)} notebooks?"):
423 raise Exit()
424 echo(
425 nbs2rich(
426 nb_paths,
427 context=export or pager,
428 )
429 )
432@app.command()
433def diff(
434 ref_base: Optional[str] = Argument(
435 None, help="Base reference (hash, branch, etc.), defaults to index"
436 ),
437 ref_remote: Optional[str] = Argument(
438 None, help="Remote reference (hash, branch, etc.), defaults to working tree"
439 ),
440 paths: List[Path] = Argument(
441 None, is_eager=True, help="Path(s) of notebook files to compare"
442 ),
443 ignore: List[str] = Option(["!*"], help="Glob expression(s) of files to ignore"),
444 export: Optional[ImgFmt] = Option(
445 None,
446 "--export",
447 "-x",
448 help="Export rich outputs as a string.",
449 ),
450 pager: bool = Option(
451 False, "--pager", "-p", help="Use pager instead of printing to terminal"
452 ),
453 verbose: bool = Option(
454 False, "--verbose", "-v", help="Increase verbosity for debugging"
455 ),
456 multiple: bool = Option(False, "--yes", "-y", help="Show multiple files"),
457 config: Optional[Path] = Option(
458 None,
459 "--config",
460 "-c",
461 is_eager=True,
462 callback=_config_callback,
463 resolve_path=True,
464 exists=True,
465 help="Get CLI options from configuration file",
466 ),
467 help: Optional[bool] = Option(
468 None,
469 "--help",
470 is_eager=True,
471 callback=_help_callback,
472 help="Show this message and exit",
473 ),
474) -> None:
475 """
476 Show differences between notebooks.
478 This is similar to `git-diff`, but in practice it is a subset of `git-diff`
479 features - only exception is that we cannot compare diffs between local files. That
480 means we can compare files that are staged with other branches, hashes, etc., or
481 compare the current directory with the current index.
482 """
483 if export is not None and pager:
484 raise BadParameter("Cannot use both pager and export output.")
485 (ref_base, ref_remote), paths = _parse_paths(ref_base, ref_remote, paths=paths)
486 diffs = get_nb_diffs(
487 ref_base=ref_base, ref_remote=ref_remote, paths=paths, verbose=verbose
488 )
489 if not diffs:
490 logger.info("No notebook diffs found. Nothing to do.")
491 raise Exit()
492 if len(diffs) > 1 and not multiple:
493 if not Confirm.ask(f"Show {len(diffs)} notebook diffs?"):
494 raise Exit()
495 echo(diffs2rich(diffs=diffs, context=export or pager))