From 3fe0c07d2679f7107de13b587e21cf31bec9d19b Mon Sep 17 00:00:00 2001
From: yh202109 <yh202109@gmail.com>
Date: Sat, 13 Jul 2024 16:59:56 -0400
Subject: [PATCH] v0.2.20

---
 docs/example_ectd.ipynb |  2 +-
 mtbp3/health/ectd.py    | 28 ++++++++++++++--------------
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/docs/example_ectd.ipynb b/docs/example_ectd.ipynb
index 092956ef..8bc7b5c4 100644
--- a/docs/example_ectd.ipynb
+++ b/docs/example_ectd.ipynb
@@ -201,7 +201,7 @@
    },
    "outputs": [],
    "source": [
-    "print(\"\\n\".join(ctoc.find_section_given_words(words=[\"REMS\",\"dsur\"], outfmt='tree', colored='nb')))"
+    "print(\"\\n\".join(ctoc.find_section_given_words(words=[\"REMS\",\"dsur\"], outfmt='tree', colored='green')))"
    ]
   },
   {
diff --git a/mtbp3/health/ectd.py b/mtbp3/health/ectd.py
index c50ba86d..270928d9 100644
--- a/mtbp3/health/ectd.py
+++ b/mtbp3/health/ectd.py
@@ -55,23 +55,23 @@ def color_output(out="", words=[], color='red'):
         elif isinstance(words, list) and words:
             assert all(isinstance(word, str) and word for word in words), "Elements in the list must be strings"
         color_dic = {'red': 31, 'green': 32, 'yellow': 33, 'blue': 34, 'magenta': 35, 'cyan': 36}
-        if color in color_dic.keys():
-            color_str = f"\x1b[{color_dic[color]}m"
+        if color not in color_dic.keys():
+            return row
         else:
-            color_str = "\x1b[31m"
+            color_str = f"\x1b[{color_dic[color]}m"
 
-        row = out
-        for word in words:
-            try:
+            row = out
+            for word in words:
+                try:
+                    start = row.lower().index(word.lower())
+                    end = start + len(word)
+                except ValueError:
+                    continue
                 start = row.lower().index(word.lower())
                 end = start + len(word)
-            except ValueError:
-                continue
-            start = row.lower().index(word.lower())
-            end = start + len(word)
-            row = row[:start] + color_str + row[start:end] + "\x1b[0m" + row[end:]
+                row = row[:start] + color_str + row[start:end] + "\x1b[0m" + row[end:]
         
-        return row
+            return row
 
     def find_section_given_words(self, words, outfmt='simple', include='up', to_right=False, colored=None):
         if isinstance(words, str) and words:
@@ -99,14 +99,14 @@ def find_section_given_words(self, words, outfmt='simple', include='up', to_righ
                     out1 = out1+out2
             out = list(set(out+out1))
 
-            if colored == "nb":
+            if colored:
                 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="blue")
+                        colored_second_part = self.color_output(second_part, words=words, color=colored)
                         out_colored.append(f"{first_part} {colored_second_part}")
                     else:
                         out_colored.append(row)