Merge pull request #60 from datarootsio/add-export-opt

Add export option

databooks/cli.py

@@ -22,7 +22,7 @@
 from databooks.logging import get_logger
 from databooks.metadata import clear_all
 from databooks.recipes import Recipe
-from databooks.tui import print_diffs, print_nbs
+from databooks.tui import ImgFmt, diffs2rich, nbs2rich
 from databooks.version import __version__
 
 logger = get_logger(__file__)
@@ -381,6 +381,12 @@ def show(
  ..., is_eager=True, help="Path(s) of notebook files with conflicts"
  ),
  ignore: List[str] = Option(["!*"], help="Glob expression(s) of files to ignore"),
+ export: Optional[ImgFmt] = Option(
+ None,
+ "--export",
+ "-x",
+ help="Export rich outputs as a string.",
+ ),
  pager: bool = Option(
  False, "--pager", "-p", help="Use pager instead of printing to terminal"
  ),
@@ -407,12 +413,18 @@ def show(
  ),
 ) -> None:
  """Show rich representation of notebook."""
+ if export is not None and pager:
+ raise BadParameter("Cannot use both pager and export output.")
  nb_paths = _check_paths(paths=paths, ignore=ignore)
  if len(nb_paths) > 1 and not multiple:
  if not Confirm.ask(f"Show {len(nb_paths)} notebooks?"):
  raise Exit()
-
- print_nbs(nb_paths, use_pager=pager)
+ echo(
+ nbs2rich(
+ nb_paths,
+ context=export or pager,
+ )
+ )
 
 
 @app.command()
@@ -427,6 +439,12 @@ def diff(
  None, is_eager=True, help="Path(s) of notebook files to compare"
  ),
  ignore: List[str] = Option(["!*"], help="Glob expression(s) of files to ignore"),
+ export: Optional[ImgFmt] = Option(
+ None,
+ "--export",
+ "-x",
+ help="Export rich outputs as a string.",
+ ),
  pager: bool = Option(
  False, "--pager", "-p", help="Use pager instead of printing to terminal"
  ),
@@ -460,6 +478,8 @@ def diff(
  means we can compare files that are staged with other branches, hashes, etc., or
  compare the current directory with the current index.
  """
+ if export is not None and pager:
+ raise BadParameter("Cannot use both pager and export output.")
  (ref_base, ref_remote), paths = _parse_paths(ref_base, ref_remote, paths=paths)
  diffs = get_nb_diffs(
  ref_base=ref_base, ref_remote=ref_remote, paths=paths, verbose=verbose
@@ -470,4 +490,4 @@ def diff(
  if len(diffs) > 1 and not multiple:
  if not Confirm.ask(f"Show {len(diffs)} notebook diffs?"):
  raise Exit()
- print_diffs(diffs=diffs, use_pager=pager)
+ echo(diffs2rich(diffs=diffs, context=export or pager))

databooks/git_utils.py

@@ -120,7 +120,6 @@ def get_repo(path: Path) -> Optional[Repo]:
  return repo
  else:
  logger.debug(f"No repo found at {path}.")
- return None
 
 
 def get_conflict_blobs(repo: Repo) -> List[ConflictFile]:

databooks/tui.py

@@ -1,8 +1,9 @@
 """Terminal user interface (TUI) helper functions and components."""
-from contextlib import nullcontext
+from contextlib import AbstractContextManager, nullcontext
 from dataclasses import asdict
+from enum import Enum
 from pathlib import Path
-from typing import List
+from typing import Any, Dict, List, Optional, Union, overload
 
 from rich.columns import Columns
 from rich.console import Console
@@ -14,30 +15,75 @@
 
 DATABOOKS_TUI = Theme({"in_count": "blue", "out_count": "orange3", "kernel": "bold"})
 
-databooks_console = Console(theme=DATABOOKS_TUI)
+ImgFmt = Enum("ImgFmt", {"html": "HTML", "svg": "SVG", "text": "TXT"})
 
 
-def print_nb(path: Path, console: Console = databooks_console) -> None:
+def nb2rich(
+ path: Path,
+ console: Console = Console(),
+) -> None:
  """Show rich representation of notebook in terminal."""
  notebook = JupyterNotebook.parse_file(path)
- console.rule(path.resolve().name)
- console.print(notebook)
+ console.print(Rule(path.resolve().name), notebook)
 
 
-def print_nbs(
+@overload
+def nbs2rich(
  paths: List[Path],
- console: Console = databooks_console,
- use_pager: bool = False,
+ *,
+ context: ImgFmt,
+ **kwargs: Any,
+) -> str:
+ ...
+
+
+@overload
+def nbs2rich(
+ paths: List[Path],
+ *,
+ context: bool,
+ **kwargs: Any,
 ) -> None:
- """Show rich representation of notebooks in terminal."""
- with console.pager(styles=True) if use_pager else nullcontext(): # type: ignore
+ ...
+
+
+def nbs2rich(
+ paths: List[Path],
+ *,
+ context: Union[ImgFmt, bool] = False,
+ export_kwargs: Optional[Dict[str, Any]] = None,
+ **console_kwargs: Any,
+) -> Optional[str]:
+ """
+ Show rich representation of notebooks in terminal.
+
+ :param paths: notebook paths to print
+ :param context: specify context - `ImgFmt` to export outputs, `True` for `pager`
+ :param export_kwargs: keyword arguments for exporting prints (as a dictionary)
+ :param console_kwargs: keyword arguments to be passed to `Console`
+ :return: console output if `context` is `ImgFmt`, else `None`
+ """
+ if "record" in console_kwargs:
+ raise ValueError(
+ "Specify `record` parameter of console via `context` argument."
+ )
+ theme = console_kwargs.pop("theme", DATABOOKS_TUI)
+ console = Console(record=isinstance(context, ImgFmt), theme=theme, **console_kwargs)
+ ctx_map: Dict[Union[ImgFmt, bool], AbstractContextManager] = {
+ True: console.pager(styles=True),
+ False: nullcontext(),
+ }
+ with ctx_map.get(context, console.capture()):
  for path in paths:
- print_nb(path, console=console)
+ nb2rich(path, console=console)
+ if isinstance(context, ImgFmt):
+ return getattr(console, f"export_{context.name}")(**(export_kwargs or {}))
 
 
-def print_diff(
+def diff2rich(
  diff: DiffContents,
- console: Console = databooks_console,
+ *,
+ console: Console = Console(),
 ) -> None:
  """Show rich representation of notebook diff in terminal."""
  a_nb, b_nb = (
@@ -62,12 +108,50 @@ def print_diff(
  console.print(cols, a_nb - b_nb)
 
 
-def print_diffs(
+@overload
+def diffs2rich(
  diffs: List[DiffContents],
- console: Console = databooks_console,
- use_pager: bool = False,
+ *,
+ context: ImgFmt,
+ **kwargs: Any,
+) -> str:
+ ...
+
+
+@overload
+def diffs2rich(
+ diffs: List[DiffContents],
+ *,
+ context: bool,
+ **kwargs: Any,
 ) -> None:
- """Show rich representation of notebook diff in terminal."""
- with console.pager(styles=True) if use_pager else nullcontext(): # type: ignore
+ ...
+
+
+def diffs2rich(
+ diffs: List[DiffContents],
+ *,
+ context: Union[ImgFmt, bool] = False,
+ export_kwargs: Optional[Dict[str, Any]] = None,
+ **console_kwargs: Any,
+) -> Optional[str]:
+ """
+ Show rich representation of notebook diff in terminal.
+
+ :param diffs: `databooks.git_utils.DiffContents` for rendering
+ :param context: specify context - `ImgFmt` to export outputs, `True` for `pager`
+ :param export_kwargs: keyword arguments for exporting prints (as a dictionary)
+ :param console_kwargs: keyword arguments to be passed to `Console`
+ :return: console output if `context` is `ImgFmt`, else `None`
+ """
+ theme = console_kwargs.pop("theme", DATABOOKS_TUI)
+ console = Console(record=isinstance(context, ImgFmt), theme=theme, **console_kwargs)
+ ctx_map: Dict[Union[ImgFmt, bool], AbstractContextManager] = {
+ True: console.pager(styles=True),
+ False: nullcontext(),
+ }
+ with ctx_map.get(context, console.capture()):
  for diff in diffs:
- print_diff(diff, console=console)
+ diff2rich(diff, console=console)
+ if isinstance(context, ImgFmt):
+ return getattr(console, f"export_{context.name}")(**(export_kwargs or {}))

docs/CLI.md

@@ -101,6 +101,7 @@ $ databooks diff [OPTIONS] [REF_BASE] [REF_REMOTE] [PATHS]...
 **Options**:
 
 * `--ignore TEXT`: Glob expression(s) of files to ignore [default: !*]
+* `-x, --export [HTML|SVG|TXT]`: Export rich outputs as a string.
 * `-p, --pager`: Use pager instead of printing to terminal [default: False]
 * `-v, --verbose`: Increase verbosity for debugging [default: False]
 * `-y, --yes`: Show multiple files [default: False]
@@ -183,6 +184,7 @@ $ databooks show [OPTIONS] PATHS...
 **Options**:
 
 * `--ignore TEXT`: Glob expression(s) of files to ignore [default: !*]
+* `-x, --export [HTML|SVG|TXT]`: Export rich outputs as a string.
 * `-p, --pager`: Use pager instead of printing to terminal [default: False]
 * `-v, --verbose`: Increase verbosity for debugging [default: False]
 * `-y, --yes`: Show multiple files [default: False]

pyproject.toml

@@ -48,6 +48,7 @@ show_error_codes = true
 disallow_untyped_calls = true
 disallow_untyped_defs = true
 ignore_missing_imports = true
+warn_no_return = false
 
 [tool.isort]
 skip_glob = ["docs/*"]

tests/test_cli.py

@@ -456,6 +456,7 @@ def test_show() -> None:
 ╭──────────────────────────────────────────────────────────────────────────────╮
 │ │
 ╰──────────────────────────────────────────────────────────────────────────────╯
+
 """
  )
 
@@ -527,6 +528,7 @@ def test_diff(tmp_path: Path) -> None:
  ╭─────────────────────────────────────╮
  <None> │ extra │
  ╰─────────────────────────────────────╯
+
 """
  )
 
@@ -556,5 +558,92 @@ def test_diff(tmp_path: Path) -> None:
  ╭─────────────────────────────────────╮
  <None> │ extra │
  ╰─────────────────────────────────────╯
+
+"""
+ )
+
+
+def test_diff_svg(tmp_path: Path) -> None:
+ """Show rich diffs of notebooks."""
+ nb_path = Path("test_conflicts_nb.ipynb")
+ notebook_1 = TestJupyterNotebook().jupyter_notebook
+ notebook_2 = TestJupyterNotebook().jupyter_notebook
+
+ notebook_1.metadata = NotebookMetadata(
+ kernelspec=dict(
+ display_name="different_kernel_display_name", name="kernel_name"
+ ),
+ field_to_remove=["Field to remove"],
+ another_field_to_remove="another field",
+ )
+
+ extra_cell = BaseCell(
+ cell_type="raw",
+ metadata=CellMetadata(random_meta=["meta"]),
+ source="extra",
+ )
+ notebook_1.cells = notebook_1.cells + [extra_cell]
+ notebook_2.nbformat += 1
+ notebook_2.nbformat_minor += 1
+
+ _ = init_repo_diff(
+ tmp_path=tmp_path,
+ filename=nb_path,
+ contents_main=notebook_1.json(),
+ contents_other=notebook_2.json(),
+ commit_message_main="Notebook from main",
+ commit_message_other="Notebook from other",
+ )
+
+ # Test passing another branch to compare
+ result = runner.invoke(app, ["diff", "other", str(tmp_path), "-x", "HTML"])
+ assert (
+ result.output
+ == """\
+<!DOCTYPE html>
+<head>
+<meta charset="UTF-8">
+<style>
+.r1 {color: #00ff00; text-decoration-color: #00ff00}
+.r2 {font-weight: bold}
+.r3 {color: #000080; text-decoration-color: #000080}
+.r4 {color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822}
+.r5 {background-color: #272822}
+body {
+ color: #000000;
+ background-color: #ffffff;
+}
+</style>
+</head>
+<html>
+<body>
+ <code>
+ <pre style="font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',\
+monospace"><span class="r1">────── </span>a/test_conflicts_nb.ipynb<span class="r1">\
+ ───────────── </span>b/test_conflicts_nb.ipynb<span class="r1"> ───────</span>
+<span class="r2"> kernel_display_name</span> <span class="r2">\
+ different_kernel_display_name</span>
+<span class="r3">In [1]:</span>
+╭──────────────────────────────────────────────────────────────────────────────╮
+│ <span class="r4">test_source</span><span class="r5"> \
+ </span> │
+╰──────────────────────────────────────────────────────────────────────────────╯
+test text
+
+<span class="r3">In [1]:</span>
+╭──────────────────────────────────────────────────────────────────────────────╮
+│ <span class="r4">test_source</span><span class="r5"> \
+ </span> │
+╰──────────────────────────────────────────────────────────────────────────────╯
+test text
+
+ ╭─────────────────────────────────────╮
+ &lt;None&gt; │ extra │
+ ╰─────────────────────────────────────╯
+</pre>
+ </code>
+</body>
+</html>
+
 """
  )

tests/test_tui.py

@@ -6,7 +6,7 @@
 
 from databooks.data_models.cell import CellMetadata, RawCell
 from databooks.data_models.notebook import JupyterNotebook, NotebookMetadata
-from databooks.tui import print_nb
+from databooks.tui import nb2rich
 from tests.test_data_models.test_notebook import TestJupyterNotebook
 
 with resources.path("tests.files", "tui-demo.ipynb") as nb_path:
@@ -193,7 +193,7 @@ def test_print_nb() -> None:
  """Print notebook from path and add rules with file name."""
  console = Console(file=io.StringIO(), width=50, legacy_windows=False)
  with resources.path("tests.files", "tui-demo.ipynb") as path:
- print_nb(path, console=console)
+ nb2rich(path, console=console)
  assert console.file.getvalue() == "\n".join(
  ("───────────────── tui-demo.ipynb ─────────────────", rich_nb)
  )