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(