slither.utils.myprettytable

 1from shutil import get_terminal_size
 2from typing import List, Dict, Union
 3
 4from prettytable import PrettyTable
 5from prettytable.colortable import ColorTable, Themes
 6
 7from slither.utils.colors import Colors
 8
 9
10class MyPrettyTable:
11    def __init__(
12        self,
13        field_names: List[str],
14        pretty_align: bool = True,
15        max_width: Union[int, None] = "max",  # Default value is "max"
16    ):
17        self._field_names = field_names
18        self._rows: List = []
19        self._options: Dict = {}
20        if pretty_align:
21            self._options["set_alignment"] = []
22            for field_name in field_names:
23                self._options["set_alignment"] += [(field_name, "l")]
24        else:
25            self._options["set_alignment"] = []
26
27        self.max_width = None
28        if max_width == "max":
29            # We use (0,0) as a fallback to detect if we are not attached to a terminal
30            # In this case, we fall back to the default behavior (i.e. printing as much as possible)
31            terminal_column = get_terminal_size((0, 0)).columns
32            if terminal_column != 0:
33                # We reduce slightly the max-width to take into account inconsistencies in terminals
34                self.max_width = terminal_column - 3
35        else:
36            self.max_width = max_width
37
38    def add_row(self, row: List[Union[str, List[str]]]) -> None:
39        self._rows.append(row)
40
41    def to_pretty_table(self) -> PrettyTable:
42        if Colors.COLORIZATION_ENABLED:
43            table = ColorTable(self._field_names, theme=Themes.OCEAN)
44        else:
45            table = PrettyTable(self._field_names)
46
47        if self.max_width is not None:
48            table.max_table_width = self.max_width
49
50        for row in self._rows:
51            table.add_row(row)
52        if len(self._options["set_alignment"]):
53            for column_header, value in self._options["set_alignment"]:
54                table.align[column_header] = value
55        return table
56
57    def to_json(self) -> Dict:
58        return {"fields_names": self._field_names, "rows": self._rows}
59
60    def __str__(self) -> str:
61        return str(self.to_pretty_table())
62
63
64# UTILITY FUNCTIONS
65
66
67def make_pretty_table(
68    headers: list, body: dict, totals: bool = False, total_header="TOTAL"
69) -> MyPrettyTable:
70    """
71    Converts a dict to a MyPrettyTable.  Dict keys are the row headers.
72    Args:
73        headers: str[] of column names
74        body: dict of row headers with a dict of the values
75        totals: bool optional add Totals row
76        total_header: str optional if totals is set to True this will override the default "TOTAL" header
77    Returns:
78        MyPrettyTable
79    """
80    table = MyPrettyTable(headers)
81    for row in body:
82        table_row = [row] + [body[row][key] for key in headers[1:]]
83        table.add_row(table_row)
84    if totals:
85        table.add_row([total_header] + [sum(body[row][key] for row in body) for key in headers[1:]])
86    return table