Coverage for databooks/data_models/rich_helpers.py: 84%

1"""Rich helpers functions for rich renderables in data models.""" 

2 

3from html.parser import HTMLParser 

4from typing import Any, List, Optional, Tuple 

5 

6from rich import box 

7from rich.console import Console, ConsoleOptions, RenderResult 

8from rich.markdown import Heading, Markdown 

9from rich.panel import Panel 

10from rich.table import Table 

11from rich.text import Text 

12 

13HtmlAttr = Tuple[str, Optional[str]] 

14 

15 

16class RichHtmlTableError(Exception): 

17 """Could not parse HTML table.""" 

18 

19 def __init__(self, msg: str = "", *args: Any): 

20 """Use class docstring as error 'prefix'.""" 

21 if self.__doc__ is None: 

22 raise ValueError("Exception docstring required - used in error message.") 

23 super().__init__(" ".join((self.__doc__, msg)), *args) 

24 

25 

26class HtmlTable(HTMLParser): 

27 """Rich table from HTML string.""" 

28 

29 def __init__(self, html: str, *args: Any, **kwargs: Any) -> None: 

30 """Initialize parser.""" 

31 super().__init__(*args, **kwargs) 

32 self.table = self.thead = self.tbody = self.body = self.th = self.td = False 

33 self.headers: List[str] = [] 

34 self.row: List[str] = [] 

35 self.rows: List[List[str]] = [] 

36 self.feed(html) 

37 

38 def handle_starttag(self, tag: str, attrs: List[HtmlAttr]) -> None: 

39 """Active tags are indicated via instance boolean properties.""" 

40 if getattr(self, tag, None): 

41 raise RichHtmlTableError(f"Already in `{tag}`.") 

42 setattr(self, tag, True) 

43 

44 def handle_endtag(self, tag: str) -> None: 

45 """Write table properties when closing tags.""" 

46 if not getattr(self, tag): 

47 raise RichHtmlTableError(f"Cannot end unopened `{tag}`.") 

48 

49 # If we are ending a row, either set a table header or row 

50 if tag == "tr": 

51 if self.thead: 

52 self.headers = self.row 

53 if self.tbody: 

54 self.rows.append(self.row) 

55 self.row = [] # restart row values 

56 setattr(self, tag, False) 

57 

58 def handle_data(self, data: str) -> None: 

59 """Append data depending on active tags.""" 

60 if self.table and (self.th or self.td): 

61 self.row.append(data) 

62 

63 def rich(self, **tbl_kwargs: Any) -> Optional[Table]: 

64 """Generate `rich` representation of table.""" 

65 if not self.rows and not self.headers: # HTML is not a table 

66 return None 

67 

68 _ncols = len(self.rows[0]) 

69 _headers = [""] * (_ncols - len(self.headers)) + self.headers 

70 if any(len(row) != _ncols for row in self.rows): 

71 raise RichHtmlTableError(f"Expected all rows to have {_ncols} columns.") 

72 

73 _box = tbl_kwargs.pop("box", box.SIMPLE_HEAVY) 

74 _row_styles = tbl_kwargs.pop("row_styles", ["on bright_black", ""]) 

75 

76 table = Table(*_headers, box=_box, row_styles=_row_styles, **tbl_kwargs) 

77 for row in self.rows: 

78 table.add_row(*row) 

79 return table 

80 

81 

82class NotebookHeading(Heading): 

83 """Preserve the notebook heading layout across Rich versions.""" 

84 

85 def __rich_console__( 

86 self, console: Console, options: ConsoleOptions 

87 ) -> RenderResult: 

88 """Center headings and draw the original double border around H1.""" 

89 text = self.text.copy() 

90 text.justify = "center" 

91 if self.tag == "h1": 

92 yield Panel(text, box=box.DOUBLE, style="markdown.h1.border") 

93 else: 

94 if self.tag == "h2": 

95 yield Text("") 

96 yield text 

97 

98 

99class NotebookMarkdown(Markdown): 

100 """Use notebook heading styles without modifying Rich's global renderer.""" 

101 

102 elements = {**Markdown.elements, "heading_open": NotebookHeading}