Coverage for databooks/tui.py: 88%

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

2from contextlib import AbstractContextManager, nullcontext 

3from dataclasses import asdict 

4from enum import Enum 

5from pathlib import Path 

6from typing import Any, Dict, List, Optional, Union, overload 

7 

8from rich.columns import Columns 

9from rich.console import Console 

10from rich.rule import Rule 

11from rich.theme import Theme 

12 

13from databooks.data_models.notebook import JupyterNotebook, NotebookMetadata 

14from databooks.git_utils import DiffContents 

15 

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

17 

18ImgFmt = Enum("ImgFmt", {"html": "HTML", "svg": "SVG", "text": "TXT"}) 

19 

20 

21def nb2rich( 

22 path: Path, 

23 console: Console = Console(), 

24) -> None: 

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

26 notebook = JupyterNotebook.parse_file(path) 

27 console.print(Rule(path.resolve().name), notebook) 

28 

29 

30@overload 

31def nbs2rich( 

32 paths: List[Path], 

33 *, 

34 context: ImgFmt, 

35 **kwargs: Any, 

36) -> str: 

37 ... 

38 

39 

40@overload 

41def nbs2rich( 

42 paths: List[Path], 

43 *, 

44 context: bool, 

45 **kwargs: Any, 

46) -> None: 

47 ... 

48 

49 

50def nbs2rich( 

51 paths: List[Path], 

52 *, 

53 context: Union[ImgFmt, bool] = False, 

54 export_kwargs: Optional[Dict[str, Any]] = None, 

55 **console_kwargs: Any, 

56) -> Optional[str]: 

57 """ 

58 Show rich representation of notebooks in terminal. 

59 

60 :param paths: notebook paths to print 

61 :param context: specify context - `ImgFmt` to export outputs, `True` for `pager` 

62 :param export_kwargs: keyword arguments for exporting prints (as a dictionary) 

63 :param console_kwargs: keyword arguments to be passed to `Console` 

64 :return: console output if `context` is `ImgFmt`, else `None` 

65 """ 

66 if "record" in console_kwargs: 

67 raise ValueError( 

68 "Specify `record` parameter of console via `context` argument." 

69 ) 

70 theme = console_kwargs.pop("theme", DATABOOKS_TUI) 

71 console = Console(record=isinstance(context, ImgFmt), theme=theme, **console_kwargs) 

72 ctx_map: Dict[Union[ImgFmt, bool], AbstractContextManager] = { 

73 True: console.pager(styles=True), 

74 False: nullcontext(), 

75 } 

76 with ctx_map.get(context, console.capture()): 

77 for path in paths: 

78 nb2rich(path, console=console) 

79 if isinstance(context, ImgFmt): 

80 return getattr(console, f"export_{context.name}")(**(export_kwargs or {})) 

81 

82 

83def diff2rich( 

84 diff: DiffContents, 

85 *, 

86 console: Console = Console(), 

87) -> None: 

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

89 a_nb, b_nb = ( 

90 JupyterNotebook.parse_raw(c) 

91 if c is not None 

92 else JupyterNotebook( 

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

94 ) 

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

96 ) 

97 cols = Columns( 

98 [ 

99 Rule( 

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

101 ) 

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

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

104 ], 

105 width=console.width // 2, 

106 padding=(0, 0), 

107 ) 

108 console.print(cols, a_nb - b_nb) 

109 

110 

111@overload 

112def diffs2rich( 

113 diffs: List[DiffContents], 

114 *, 

115 context: ImgFmt, 

116 **kwargs: Any, 

117) -> str: 

118 ... 

119 

120 

121@overload 

122def diffs2rich( 

123 diffs: List[DiffContents], 

124 *, 

125 context: bool, 

126 **kwargs: Any, 

127) -> None: 

128 ... 

129 

130 

131def diffs2rich( 

132 diffs: List[DiffContents], 

133 *, 

134 context: Union[ImgFmt, bool] = False, 

135 export_kwargs: Optional[Dict[str, Any]] = None, 

136 **console_kwargs: Any, 

137) -> Optional[str]: 

138 """ 

139 Show rich representation of notebook diff in terminal. 

140 

141 :param diffs: `databooks.git_utils.DiffContents` for rendering 

142 :param context: specify context - `ImgFmt` to export outputs, `True` for `pager` 

143 :param export_kwargs: keyword arguments for exporting prints (as a dictionary) 

144 :param console_kwargs: keyword arguments to be passed to `Console` 

145 :return: console output if `context` is `ImgFmt`, else `None` 

146 """ 

147 theme = console_kwargs.pop("theme", DATABOOKS_TUI) 

148 console = Console(record=isinstance(context, ImgFmt), theme=theme, **console_kwargs) 

149 ctx_map: Dict[Union[ImgFmt, bool], AbstractContextManager] = { 

150 True: console.pager(styles=True), 

151 False: nullcontext(), 

152 } 

153 with ctx_map.get(context, console.capture()): 

154 for diff in diffs: 

155 diff2rich(diff, console=console) 

156 if isinstance(context, ImgFmt): 

157 return getattr(console, f"export_{context.name}")(**(export_kwargs or {}))