Coverage for databooks/tui.py: 93%

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

2 

3from contextlib import AbstractContextManager, nullcontext 

4from dataclasses import asdict 

5from enum import Enum 

6from io import StringIO 

7from pathlib import Path 

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

9 

10from rich.columns import Columns 

11from rich.console import Console 

12from rich.rule import Rule 

13from rich.theme import Theme 

14 

15from databooks.data_models.notebook import JupyterNotebook, NotebookMetadata 

16from databooks.git_utils import DiffContents 

17 

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

19 

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

21 

22 

23def nb2rich( 

24 path: Path, 

25 console: Console = Console(), 

26) -> None: 

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

28 notebook = JupyterNotebook.parse_file(path) 

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

30 

31 

32@overload 

33def nbs2rich( 

34 paths: List[Path], 

35 *, 

36 context: ImgFmt, 

37 **kwargs: Any, 

38) -> str: ... 

39 

40 

41@overload 

42def nbs2rich( 

43 paths: List[Path], 

44 *, 

45 context: bool, 

46 **kwargs: Any, 

47) -> None: ... 

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 if isinstance(context, ImgFmt): 

72 console_kwargs["file"] = StringIO() 

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

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

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

76 False: nullcontext(), 

77 } 

78 with ctx_map.get(context, nullcontext()): 

79 for path in paths: 

80 nb2rich(path, console=console) 

81 if isinstance(context, ImgFmt): 

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

83 

84 

85def diff2rich( 

86 diff: DiffContents, 

87 *, 

88 console: Console = Console(), 

89) -> None: 

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

91 a_nb, b_nb = ( 

92 ( 

93 JupyterNotebook.model_validate_json(c) 

94 if c is not None 

95 else JupyterNotebook( 

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

97 ) 

98 ) 

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

100 ) 

101 cols = Columns( 

102 [ 

103 Rule( 

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

105 ) 

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

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

108 ], 

109 width=console.width // 2, 

110 padding=(0, 0), 

111 ) 

112 console.print(cols, a_nb - b_nb) 

113 

114 

115@overload 

116def diffs2rich( 

117 diffs: List[DiffContents], 

118 *, 

119 context: ImgFmt, 

120 **kwargs: Any, 

121) -> str: ... 

122 

123 

124@overload 

125def diffs2rich( 

126 diffs: List[DiffContents], 

127 *, 

128 context: bool, 

129 **kwargs: Any, 

130) -> None: ... 

131 

132 

133def diffs2rich( 

134 diffs: List[DiffContents], 

135 *, 

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

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

138 **console_kwargs: Any, 

139) -> Optional[str]: 

140 """ 

141 Show rich representation of notebook diff in terminal. 

142 

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

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

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

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

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

148 """ 

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

150 if isinstance(context, ImgFmt): 

151 console_kwargs["file"] = StringIO() 

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

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

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

155 False: nullcontext(), 

156 } 

157 with ctx_map.get(context, nullcontext()): 

158 for diff in diffs: 

159 diff2rich(diff, console=console) 

160 if isinstance(context, ImgFmt): 

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