slither: utils: respect colorization state when printing tables

Closes crytic#2229
elopez · Feb 16, 2024 · b66b3e0 · b66b3e0
1 parent e876d61
commit b66b3e0
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/slither/utils/myprettytable.py b/slither/utils/myprettytable.py
@@ -1,7 +1,10 @@
 from typing import List, Dict, Union
 
+from prettytable import PrettyTable
 from prettytable.colortable import ColorTable, Themes
 
+from slither.utils.colors import Colors
+
 
 class MyPrettyTable:
     def __init__(self, field_names: List[str], pretty_align: bool = True):  # TODO: True by default?
@@ -19,8 +22,12 @@ def __init__(self, field_names: List[str], pretty_align: bool = True):  # TODO:
     def add_row(self, row: List[Union[str, List[str]]]) -> None:
         self._rows.append(row)
 
-    def to_pretty_table(self) -> ColorTable:
-        table = ColorTable(self._field_names, theme=Themes.OCEAN)
+    def to_pretty_table(self) -> PrettyTable:
+        if Colors.COLORIZATION_ENABLED:
+            table = ColorTable(self._field_names, theme=Themes.OCEAN)
+        else:
+            table = PrettyTable(self._field_names)
+
         for row in self._rows:
             table.add_row(row)
         if len(self._options["set_alignment"]):