Coverage for databooks/data_models/cell.py: 94%

1"""Data models - Cells and components.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Union, cast 

6 

7from pydantic import PositiveInt, RootModel, field_validator 

8from rich.console import Console, ConsoleOptions, ConsoleRenderable, RenderResult 

9from rich.panel import Panel 

10from rich.syntax import Syntax 

11from rich.text import Text 

12 

13from databooks.data_models.base import DatabooksBase 

14from databooks.data_models.rich_helpers import ( 

15 HtmlTable, 

16 NotebookMarkdown, 

17 RichHtmlTableError, 

18) 

19from databooks.logging import get_logger 

20 

21logger = get_logger(__file__) 

22 

23 

24class CellMetadata(DatabooksBase): 

25 """Cell metadata. Empty by default but can accept extra fields.""" 

26 

27 

28class BaseCell(DatabooksBase): 

29 """ 

30 Jupyter notebook cells. 

31 

32 Fields `outputs` and `execution_count` are not included since they should only be 

33 present in code cells - thus are treated as extra fields. 

34 """ 

35 

36 metadata: CellMetadata 

37 source: Union[List[str], str] 

38 cell_type: str 

39 

40 def __hash__(self) -> int: 

41 """Cells must be hashable for `difflib.SequenceMatcher`.""" 

42 return hash( 

43 (type(self),) + tuple(v) if isinstance(v, list) else v 

44 for v in self.__dict__.values() 

45 ) 

46 

47 def remove_fields( 

48 self, fields: Iterable[str] = (), missing_ok: bool = True, **kwargs: Any 

49 ) -> None: 

50 """ 

51 Remove cell fields. 

52 

53 Similar to `databooks.data_models.base.remove_fields`, but will ignore required 

54 fields for cell type. 

55 """ 

56 # Ignore required `BaseCell` fields 

57 cell_fields = BaseCell.model_fields # required fields 

58 if any(field in fields for field in cell_fields): 

59 logger.debug( 

60 "Ignoring removal of required fields " 

61 + str([f for f in fields if f in cell_fields]) 

62 + f" in `{type(self).__name__}`." 

63 ) 

64 fields = [f for f in fields if f not in cell_fields] 

65 

66 super(BaseCell, self).remove_fields(fields, missing_ok=missing_ok) 

67 

68 if self.cell_type == "code": 

69 self.outputs: CellOutputs = ( 

70 CellOutputs([]) if "outputs" not in dict(self) else self.outputs 

71 ) 

72 self.execution_count: Optional[PositiveInt] = ( 

73 None if "execution_count" not in dict(self) else self.execution_count 

74 ) 

75 

76 def clear_fields( 

77 self, 

78 *, 

79 cell_metadata_keep: Optional[Sequence[str]] = None, 

80 cell_metadata_remove: Optional[Sequence[str]] = None, 

81 cell_remove_fields: Sequence[str] = (), 

82 ) -> None: 

83 """ 

84 Clear cell metadata, execution count, outputs or other desired fields (id, ...). 

85 

86 You can also specify metadata to keep or remove from the `metadata` property of 

87 `databooks.data_models.cell.BaseCell`. 

88 :param cell_metadata_keep: Metadata values to keep - simply pass an empty 

89 sequence (i.e.: `()`) to remove all extra fields. 

90 :param cell_metadata_remove: Metadata values to remove 

91 :param cell_remove_fields: Fields to remove from cell 

92 :return: 

93 """ 

94 nargs = sum((cell_metadata_keep is not None, cell_metadata_remove is not None)) 

95 if nargs != 1: 

96 raise ValueError( 

97 "Exactly one of `cell_metadata_keep` or `cell_metadata_remove` must" 

98 f" be passed, got {nargs} arguments." 

99 ) 

100 

101 if cell_metadata_keep is not None: 

102 cell_metadata_remove = tuple( 

103 field for field, _ in self.metadata if field not in cell_metadata_keep 

104 ) 

105 # Non-None here: exactly one of keep/remove is given and keep became remove 

106 self.metadata.remove_fields(cast(Sequence[str], cell_metadata_remove)) 

107 

108 self.remove_fields(fields=cell_remove_fields, missing_ok=True) 

109 

110 

111class CellStreamOutput(DatabooksBase): 

112 """Cell output of type `stream`.""" 

113 

114 output_type: str 

115 name: str 

116 text: List[str] 

117 

118 def __rich__( 

119 self, 

120 ) -> ConsoleRenderable: 

121 """Rich display of cell stream outputs.""" 

122 return Text("".join(self.text)) 

123 

124 @field_validator("output_type") 

125 @classmethod 

126 def output_type_must_be_stream(cls, v: str) -> str: 

127 """Check if stream has `stream` type.""" 

128 if v != "stream": 

129 raise ValueError(f"Invalid output type. Expected `stream`, got {v}.") 

130 return v 

131 

132 @field_validator("name") 

133 @classmethod 

134 def stream_name_must_match(cls, v: str) -> str: 

135 """Check if stream name is either `stdout` or `stderr`.""" 

136 valid_names = ("stdout", "stderr") 

137 if v not in valid_names: 

138 raise ValueError( 

139 f"Invalid stream name. Expected one of {valid_names}, got {v}." 

140 ) 

141 return v 

142 

143 

144class CellDisplayDataOutput(DatabooksBase): 

145 """Cell output of type `display_data`.""" 

146 

147 output_type: str 

148 data: Dict[str, Any] 

149 metadata: Dict[str, Any] 

150 

151 @property 

152 def rich_output(self) -> Sequence[ConsoleRenderable]: 

153 """Dynamically compute the rich output - also in `CellExecuteResultOutput`.""" 

154 

155 def _try_parse_html(s: str) -> Optional[ConsoleRenderable]: 

156 """Try to parse HTML table, return `None` if any errors are raised.""" 

157 try: 

158 return HtmlTable("".join(s)).rich() 

159 except RichHtmlTableError: 

160 logger.debug("Could not generate rich HTML table.") 

161 return None 

162 

163 mime_func: Dict[str, Callable[[str], Optional[ConsoleRenderable]]] = { 

164 "text/html": lambda s: _try_parse_html(s), 

165 "text/plain": lambda s: Text("".join(s)), 

166 } 

167 

168 def unsupported(s: str) -> None: 

169 return None 

170 

171 _rich = { 

172 mime: mime_func.get(mime, unsupported)(content) # try to render element 

173 for mime, content in self.data.items() 

174 } 

175 return [ 

176 *[ 

177 Text(f"<✨Rich✨ `{mime}` not available 😢>") 

178 for mime, renderable in _rich.items() 

179 if renderable is None 

180 ], 

181 next(renderable for renderable in _rich.values() if renderable is not None), 

182 ] 

183 

184 def __rich_console__( 

185 self, console: Console, options: ConsoleOptions 

186 ) -> RenderResult: 

187 """Rich display of data display outputs.""" 

188 yield from self.rich_output 

189 

190 @field_validator("output_type") 

191 @classmethod 

192 def output_type_must_match(cls, v: str) -> str: 

193 """Check if stream has `display_data` type.""" 

194 if v != "display_data": 

195 raise ValueError(f"Invalid output type. Expected `display_data`, got {v}.") 

196 return v 

197 

198 

199class CellExecuteResultOutput(CellDisplayDataOutput): 

200 """Cell output of type `execute_result`.""" 

201 

202 execution_count: PositiveInt 

203 

204 def __rich_console__( 

205 self, console: Console, options: ConsoleOptions 

206 ) -> RenderResult: 

207 """Rich display of executed cell outputs.""" 

208 yield Text(f"Out [{self.execution_count or ' '}]:", style="out_count") 

209 yield from self.rich_output 

210 

211 @field_validator("output_type") 

212 @classmethod 

213 def output_type_must_match(cls, v: str) -> str: 

214 """Check if stream has `execute_result` type.""" 

215 if v != "execute_result": 

216 raise ValueError( 

217 f"Invalid output type. Expected `execute_result`, got {v}." 

218 ) 

219 return v 

220 

221 

222class CellErrorOutput(DatabooksBase): 

223 """Cell output of type `error`.""" 

224 

225 output_type: str 

226 ename: str 

227 evalue: str 

228 traceback: List[str] 

229 

230 def __rich__( 

231 self, 

232 ) -> ConsoleRenderable: 

233 """Rich display of error outputs.""" 

234 return Text.from_ansi("\n".join(self.traceback).rstrip("\n")) 

235 

236 @field_validator("output_type") 

237 @classmethod 

238 def output_type_must_match(cls, v: str) -> str: 

239 """Check if stream has `error` type.""" 

240 if v != "error": 

241 raise ValueError(f"Invalid output type. Expected `error`, got {v}.") 

242 return v 

243 

244 

245CellOutputType = Union[ 

246 CellStreamOutput, CellDisplayDataOutput, CellExecuteResultOutput, CellErrorOutput 

247] 

248 

249 

250class CellOutputs(RootModel): 

251 """Outputs of notebook code cells.""" 

252 

253 root: List[CellOutputType] 

254 

255 def __rich_console__( 

256 self, console: Console, options: ConsoleOptions 

257 ) -> RenderResult: 

258 """Rich display of code cell outputs.""" 

259 yield from self.values 

260 

261 @property 

262 def values( 

263 self, 

264 ) -> List[CellOutputType]: 

265 """Alias `root` with outputs for easy referencing.""" 

266 return self.root 

267 

268 

269class CodeCell(BaseCell): 

270 """Cell of type `code` - defined for rich displaying in terminal.""" 

271 

272 outputs: CellOutputs 

273 cell_type: str = "code" 

274 

275 def __rich_console__( 

276 self, console: Console, options: ConsoleOptions 

277 ) -> RenderResult: 

278 """Rich display of code cells.""" 

279 yield Text(f"In [{self.execution_count or ' '}]:", style="in_count") 

280 yield Panel( 

281 Syntax( 

282 "".join(self.source) if isinstance(self.source, list) else self.source, 

283 getattr(self.metadata, "lang", "text"), 

284 ) 

285 ) 

286 yield self.outputs 

287 

288 @field_validator("cell_type") 

289 @classmethod 

290 def cell_has_code_type(cls, v: str) -> str: 

291 """Extract the list values from the __root__ attribute of `CellOutputs`.""" 

292 if v != "code": 

293 raise ValueError(f"Expected code of type `code`, got `{v}`.") 

294 return v 

295 

296 

297class MarkdownCell(BaseCell): 

298 """Cell of type `markdown` - defined for rich displaying in terminal.""" 

299 

300 cell_type: str = "markdown" 

301 

302 def __rich__( 

303 self, 

304 ) -> ConsoleRenderable: 

305 """Rich display of markdown cells.""" 

306 return Panel(NotebookMarkdown("".join(self.source))) 

307 

308 @field_validator("cell_type") 

309 @classmethod 

310 def cell_has_md_type(cls, v: str) -> str: 

311 """Extract the list values from the __root__ attribute of `CellOutputs`.""" 

312 if v != "markdown": 

313 raise ValueError(f"Expected code of type `markdown`, got {v}.") 

314 return v 

315 

316 

317class RawCell(BaseCell): 

318 """Cell of type `raw` - defined for rich displaying in terminal.""" 

319 

320 cell_type: str = "raw" 

321 

322 def __rich__( 

323 self, 

324 ) -> ConsoleRenderable: 

325 """Rich display of raw cells.""" 

326 return Panel(Text("".join(self.source))) 

327 

328 @field_validator("cell_type") 

329 @classmethod 

330 def cell_has_md_type(cls, v: str) -> str: 

331 """Extract the list values from the __root__ attribute of `CellOutputs`.""" 

332 if v != "raw": 

333 raise ValueError(f"Expected code of type `raw`, got {v}.") 

334 return v