From f3b69f2a714843892605c70449c7db89bbb6b613 Mon Sep 17 00:00:00 2001
From: Maximilian Bosch <maximilian@mbosch.me>
Date: Sun, 19 Apr 2026 12:17:18 +0200
Subject: [PATCH] nixos/test-driver: colorized warnings

---
 .../test-driver/src/test_driver/__init__.py   | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/nixos/lib/test-driver/src/test_driver/__init__.py b/nixos/lib/test-driver/src/test_driver/__init__.py
index 1e5a062b18c8..95f25136701f 100644
--- a/nixos/lib/test-driver/src/test_driver/__init__.py
+++ b/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(