v0.2.20

yh202109 · Jul 13, 2024 · ff5f2e8 · ff5f2e8
1 parent 01f3cc0
commit ff5f2e8
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 13 deletions.
diff --git a/docs/example_ectd2.ipynb b/docs/example_ectd2.ipynb
@@ -202,7 +202,7 @@
    },
    "outputs": [],
    "source": [
-    "print(\"\\n\".join(ctoc.find_section_given_words(words=[\"REMS\",\"dsur\"], outfmt='tree')))"
+    "print(\"\\n\".join(ctoc.find_section_given_words(words=[\"REMS\",\"dsur\"], outfmt='tree', colored='nb')))"
    ]
   },
   {

diff --git a/mtbp3/health/ectd.py b/mtbp3/health/ectd.py
@@ -68,7 +68,7 @@ def color_output(out="", words=[], color='red'):
 
         return row
 
-    def find_section_given_words(self, words, outfmt='simple', include='up', to_right=False):
+    def find_section_given_words(self, words, outfmt='simple', include='up', to_right=False, colored=None):
         if isinstance(words, str) and words:
             words = [words]
         elif isinstance(words, list) and len(words)>0:
@@ -93,18 +93,21 @@ def find_section_given_words(self, words, outfmt='simple', include='up', to_righ
                     out2 = [item for item in self.ctoc if item.startswith(out0_str1 + " ")]
                     out1 = out1+out2
             out = list(set(out+out1))
-            out_colored=[]
-            for row in out: 
-                split_row = row.split(' ', 1)
-                if len(split_row) > 1:
-                    first_part = split_row[0]
-                    second_part = split_row[1]
-                    colored_second_part = self.color_output(second_part, words=words, color="red")
-                    out_colored.append(f"{first_part} {colored_second_part}")
-                else:
-                    out_colored.append(row)
 
-            out_tree = ListTree(lst=out_colored, infmt='dotspace')
+            if colored == "nb":
+                out_colored=[]
+                for row in out: 
+                    split_row = row.split(' ', 1)
+                    if len(split_row) > 1:
+                        first_part = split_row[0]
+                        second_part = split_row[1]
+                        colored_second_part = self.color_output(second_part, words=words, color="red")
+                        out_colored.append(f"{first_part} {colored_second_part}")
+                    else:
+                        out_colored.append(row)
+                out_tree = ListTree(lst=out_colored, infmt='dotspace')
+            else:
+                out_tree = ListTree(lst=out, infmt='dotspace')
             return out_tree.list_tree(to_right=to_right)
         else:
             raise ValueError("Invalid value for outfmt. Supported values are 'simple' and 'tree'.")