From 4a1486b1db77d01cc397d47f48fe10a838403fe2 Mon Sep 17 00:00:00 2001
From: Florent Kermarrec <florent@enjoy-digital.fr>
Date: Fri, 3 Nov 2023 11:06:41 +0100
Subject: [PATCH] gen/fhdl/hierarchy: Add with_colors parameters to allow
 enabling/disabling colors.

---
 litex/gen/fhdl/hierarchy.py | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/litex/gen/fhdl/hierarchy.py b/litex/gen/fhdl/hierarchy.py
index eb8ffde99..17ff922dc 100644
--- a/litex/gen/fhdl/hierarchy.py
+++ b/litex/gen/fhdl/hierarchy.py
@@ -14,11 +14,15 @@ class LiteXHierarchyExplorer:
     tree_ident      = "│    "
     tree_entry      = "└─── "
 
-    def __init__(self, top, depth=None):
-        self.top   = top
-        self.depth = depth
+    def __init__(self, top, depth=None, with_colors=True):
+        self.top         = top
+        self.depth       = depth
+        self.with_colors = with_colors
 
-    def get_tree(self, module, ident=0, with_modules=True, with_instances=True):
+    def _colorer(self, s, color="bright"):
+        return colorer(s=s, color=color, enable=self.with_colors)
+
+    def get_tree(self, module, ident=0, with_modules=True, with_instances=True, with_colors=True):
         r = ""
         names = set()
         names.add(None)
@@ -31,7 +35,7 @@ class LiteXHierarchyExplorer:
                     n += 1
             names.add(name)
             if with_modules:
-                r += f"{self.tree_ident*ident}{self.tree_entry}{colorer(name, 'cyan')} ({mod.__class__.__name__})\n"
+                r += f"{self.tree_ident*ident}{self.tree_entry}{self._colorer(name, 'cyan')} ({mod.__class__.__name__})\n"
             if (self.depth is None) or (ident < self.depth):
                 r += self.get_tree(mod, ident + 1)
 
@@ -44,13 +48,13 @@ class LiteXHierarchyExplorer:
                         if s in v._fragment.specials:
                             show = False
                     if show:
-                        r +=  f"{self.tree_ident*ident}{self.tree_entry}{colorer(f'[{s.of}]', 'yellow')}\n"
+                        r +=  f"{self.tree_ident*ident}{self.tree_entry}{self._colorer(f'[{s.of}]', 'yellow')}\n"
         return r
 
     def __repr__(self):
         r = "\n"
-        r += f"{colorer(self.top.__class__.__name__, 'underline')}\n"
+        r += f"{self._colorer(self.top.__class__.__name__, 'underline')}\n"
         r += self.get_tree(self.top)
-        r += f"{colorer('* ', 'cyan')}: Generated name.\n"
-        r += f"{colorer('[]', 'yellow')}: BlackBox.\n"
+        r += f"{self._colorer('* ', 'cyan')}: Generated name.\n"
+        r += f"{self._colorer('[]', 'yellow')}: BlackBox.\n"
         return r