mavsdk_tests: import naming, easier coloring

test/mavsdk_tests/logger_helper.py

@@ -3,50 +3,52 @@
 import re
 import sys
 import os
+from enum import Enum
+from functools import lru_cache
 
 
-def supports_color():
-    """Returns True if the running system's terminal supports color.
-
-    From https://stackoverflow.com/a/22254892/8548472
-    """
-    plat = sys.platform
-    supported_platform = plat != 'Pocket PC' and (plat != 'win32' or
-                                                  'ANSICON' in os.environ)
-    # isatty is not always implemented, #6223.
-    is_a_tty = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()
-    return supported_platform and is_a_tty
+class color(Enum):
+    PURPLE = '\033[95m'
+    CYAN = '\033[96m'
+    DARKCYAN = '\033[36m'
+    BLUE = '\033[94m'
+    GREEN = '\033[92m'
+    YELLOW = '\033[93m'
+    RED = '\033[91m'
+    BOLD = '\033[1m'
+    UNDERLINE = '\033[4m'
+    RESET = '\033[0m'
 
 
-if supports_color():
-    class color:
-        PURPLE = '\033[95m'
-        CYAN = '\033[96m'
-        DARKCYAN = '\033[36m'
-        BLUE = '\033[94m'
-        GREEN = '\033[92m'
-        YELLOW = '\033[93m'
-        RED = '\033[91m'
-        BOLD = '\033[1m'
-        UNDERLINE = '\033[4m'
-        END = '\033[0m'
-else:
-    class color:
-        PURPLE = ''
-        CYAN = ''
-        DARKCYAN = ''
-        BLUE = ''
-        GREEN = ''
-        YELLOW = ''
-        RED = ''
-        BOLD = ''
-        UNDERLINE = ''
-        END = ''
+def colorize(text: str, c: color) -> str:
+    if _supports_color():
+        return c.value + text + color.RESET.value
+    else:
+        return text
 
 
-def remove_color(text):
+def maybe_strip_color(text: str) -> str:
     """Remove ANSI and xterm256 color codes.
 
     From https://stackoverflow.com/a/30500866/8548472
     """
-    return re.sub(r'\x1b(\[.*?[@-~]|\].*?(\x07|\x1b\\))', '', text)
+    if not _supports_color():
+        return re.sub(r'\x1b(\[.*?[@-~]|\].*?(\x07|\x1b\\))', '', text)
+    else:
+        return text
+
+
+@lru_cache
+def _supports_color() -> bool:
+    """Returns True if the running system's terminal supports color.
+
+    From https://stackoverflow.com/a/22254892/8548472
+    """
+    supported_platform = \
+        (sys.platform != 'Pocket PC') and \
+        (sys.platform != 'win32' or 'ANSICON' in os.environ)
+
+    # isatty is not always implemented.
+    is_a_tty = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()
+
+    return supported_platform and is_a_tty