Coverage for databooks/tui.py: 100%

1"""Terminal user interface (TUI) helper functions and components.""" 

2from contextlib import nullcontext 

3from dataclasses import asdict 

4from pathlib import Path 

5from typing import List 

6 

7from rich.columns import Columns 

8from rich.console import Console 

9from rich.rule import Rule 

10from rich.theme import Theme 

11 

12from databooks.data_models.notebook import JupyterNotebook, NotebookMetadata 

13from databooks.git_utils import DiffContents 

14 

15DATABOOKS_TUI = Theme({"in_count": "blue", "out_count": "orange3", "kernel": "bold"}) 

16 

17databooks_console = Console(theme=DATABOOKS_TUI) 

18 

19 

20def print_nb(path: Path, console: Console = databooks_console) -> None: 

21 """Show rich representation of notebook in terminal.""" 

22 notebook = JupyterNotebook.parse_file(path) 

23 console.rule(path.resolve().name) 

24 console.print(notebook) 

25 

26 

27def print_nbs( 

28 paths: List[Path], 

29 console: Console = databooks_console, 

30 use_pager: bool = False, 

31) -> None: 

32 """Show rich representation of notebooks in terminal.""" 

33 with console.pager(styles=True) if use_pager else nullcontext(): # type: ignore 

34 for path in paths: 

35 print_nb(path, console=console) 

36 

37 

38def print_diff( 

39 diff: DiffContents, 

40 console: Console = databooks_console, 

41) -> None: 

42 """Show rich representation of notebook diff in terminal.""" 

43 a_nb, b_nb = ( 

44 JupyterNotebook.parse_raw(c) 

45 if c is not None 

46 else JupyterNotebook( 

47 nbformat=0, nbformat_minor=0, metadata=NotebookMetadata(), cells=[] 

48 ) 

49 for c in (diff.a.contents, diff.b.contents) 

50 ) 

51 cols = Columns( 

52 [ 

53 Rule( 

54 f"{ab}/{c['path'].resolve().name if c['path'] is not None else 'null'}" 

55 ) 

56 for ab, c in asdict(diff).items() 

57 if ab in ("a", "b") 

58 ], 

59 width=console.width // 2, 

60 padding=(0, 0), 

61 ) 

62 console.print(cols, a_nb - b_nb) 

63 

64 

65def print_diffs( 

66 diffs: List[DiffContents], 

67 console: Console = databooks_console, 

68 use_pager: bool = False, 

69) -> None: 

70 """Show rich representation of notebook diff in terminal.""" 

71 with console.pager(styles=True) if use_pager else nullcontext(): # type: ignore 

72 for diff in diffs: 

73 print_diff(diff, console=console)