nixos/test-driver: colorized warnings

nixos/lib/test-driver/src/test_driver/__init__.py

@@ -6,6 +6,7 @@ import warnings
 from pathlib import Path
 
 import ptpython.ipython
+from colorama import Fore, Style
 
 from test_driver.debug import Debug, DebugAbstract, DebugNop
 from test_driver.driver import Driver, DriverConfiguration, load_driver_configuration
@@ -55,6 +56,29 @@ def writeable_dir(arg: str) -> Path:
     return path
 
 
+def formatwarning(
+    message: Warning | str,
+    category: type[Warning],
+    filename: str,
+    lineno: int,
+    line: str | None = None,
+) -> str:
+    return (
+        Style.BRIGHT
+        + Fore.YELLOW
+        + f"??? Warning ({category.__name__}): "  # ty: ignore[unsupported-operator]
+        + Style.NORMAL
+        + str(message)
+        + "\n"
+        + f'    File "{filename}", line {lineno}\n'
+        + (f"      {line}\n" if line is not None else "")
+        + Style.RESET_ALL
+    )
+
+
+warnings.formatwarning = formatwarning  # ty:ignore[invalid-assignment]
+
+
 def main() -> None:
     arg_parser = argparse.ArgumentParser(prog="nixos-test-driver")
     arg_parser.add_argument(