From a6c469baae717b08cd110168eeac8dabbf21f23b Mon Sep 17 00:00:00 2001
From: Kevin Weiss <weiss.kevin604@gmail.com>
Date: Tue, 18 Nov 2025 13:33:18 +0100
Subject: [PATCH] feat: Allow lob_print to have color

This allows an optional kwarg color when printing.  It does not need any dependencies and should fallback safely.
---
 src/lob_hlpr/hlpr.py   | 70 +++++++++++++++++++++++++++++++++++++++++-
 tests/test_lob_hlpr.py |  5 +++
 2 files changed, 74 insertions(+), 1 deletion(-)

diff --git a/src/lob_hlpr/hlpr.py b/src/lob_hlpr/hlpr.py
index d83a37a..1851499 100644
--- a/src/lob_hlpr/hlpr.py
+++ b/src/lob_hlpr/hlpr.py
@@ -9,6 +9,38 @@
 from lob_hlpr.lib_types import FirmwareID
 
 
+def enable_windows_ansi_support():  # pragma: no cover
+    """Try to enable ANSI escape sequence support on Windows.
+
+    Works on Windows 10+.
+    """
+    if os.name != "nt":
+        return True  # Non-Windows always supports ANSI
+
+    try:
+        import ctypes
+
+        kernel32 = ctypes.windll.kernel32
+
+        # Enable Virtual Terminal Processing
+        handle = kernel32.GetStdHandle(-11)  # STD_OUTPUT_HANDLE = -11
+        mode = ctypes.c_uint32()
+        if kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
+            new_mode = (
+                mode.value | 0x0004
+            )  # ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
+            if kernel32.SetConsoleMode(handle, new_mode):
+                return True
+    except Exception:
+        pass
+
+    return False
+
+
+# Determine if ANSI colors will work
+_USE_COLOR = enable_windows_ansi_support()
+
+
 class LobHlpr:
     """Helper functions for Lobaro tools."""
 
@@ -54,6 +86,31 @@ def sn_vid_pid_to_regex(
             sn = re.escape(sn)
         return f"VID:PID={vid or '.*'}:{pid or '.*'}.+SER={sn}"
 
+    @staticmethod
+    def _print_color(*args, color=None, **kwargs):
+        """Print with color if supported."""
+        # ANSI color codes
+        RESET = "\033[0m"
+        RED = "\033[31m"
+        GREEN = "\033[32m"
+        YELLOW = "\033[33m"
+        if color is None or not _USE_COLOR:
+            print(*args, flush=True, **kwargs)
+            return
+        text = kwargs.get("sep", " ").join(str(a) for a in args)
+        color = color.lower()
+        code = None
+        if "red" == color:
+            code = RED
+        elif "green" == color:
+            code = GREEN
+        elif "yellow" == color:
+            code = YELLOW
+        if code is not None:
+            print(f"{code}{text}{RESET}", flush=True, **kwargs)
+        else:
+            print(text, flush=True, **kwargs)
+
     @staticmethod
     def lob_print(log_path: str, *args, **kwargs):
         """Print to the console and log to a file.
@@ -61,8 +118,19 @@ def lob_print(log_path: str, *args, **kwargs):
         The log file is rotated when it reaches 256MB and the last two
         log files are kept. This can write all log messages to the log file
         only if the log handlers are set (i.e. basicConfig loglevel is Debug).
+
+        Args:
+            log_path: The path to the log file.
+            *args: Arguments to print.
+            **kwargs: Additional keyword arguments.
+                color (str, optional): If provided, prints the message in color
+                    to the console. Supported values are "red", "green", and "yellow".
+                    If colors are not supported by the terminal, output will be
+                    uncolored.
         """
-        print(*args, flush=True, **kwargs)
+        color = kwargs.pop("color", None)
+        LobHlpr._print_color(*args, color=color, **kwargs)
+
         # get the directory from the log_path
         log_dir = os.path.dirname(log_path)
         os.makedirs(log_dir, exist_ok=True)
diff --git a/tests/test_lob_hlpr.py b/tests/test_lob_hlpr.py
index 47a16dd..8d6570c 100644
--- a/tests/test_lob_hlpr.py
+++ b/tests/test_lob_hlpr.py
@@ -138,12 +138,17 @@ def test_log_print_passes(tmp_path, capsys):
     test_file = tmp_path / "test.log"
     hlp.lob_print(str(test_file), "Test message")
     hlp.lob_print(str(test_file), "Another test message")
+    hlp.lob_print(str(test_file), "red message", color="red")
+    hlp.lob_print(str(test_file), "yellow message", color="yellow")
+    hlp.lob_print(str(test_file), "green message", color="green")
+    hlp.lob_print(str(test_file), "normal message", color="doesn't matter")
     test_logger.info("Will show in logs after lob_print")
     captured = capsys.readouterr()
     assert "Test message" in captured.out
     # Check that only one "Another test message" is in the output
     assert captured.out.count("Another test message") == 1
     assert test_file.exists()
+    assert captured.out.count("red message") == 1
     with open(test_file) as f:
         log_content = f.read()
         assert "Test message" in log_content