testing/kunit/kunit_printer.py

0001 #!/usr/bin/env python3
0002 # SPDX-License-Identifier: GPL-2.0
0003 #
0004 # Utilities for printing and coloring output.
0005 #
0006 # Copyright (C) 2022, Google LLC.
0007 # Author: Daniel Latypov <dlatypov@google.com>
0008
0009 import datetime
0010 import sys
0011 import typing
0012
0013 _RESET = '\033[0;0m'
0014
0015 class Printer:
0016     """Wraps a file object, providing utilities for coloring output, etc."""
0017
0018     def __init__(self, output: typing.IO):
0019         self._output = output
0020         self._use_color = output.isatty()
0021
0022     def print(self, message: str) -> None:
0023         print(message, file=self._output)
0024
0025     def print_with_timestamp(self, message: str) -> None:
0026         ts = datetime.datetime.now().strftime('%H:%M:%S')
0027         self.print(f'[{ts}] {message}')
0028
0029     def _color(self, code: str, text: str) -> str:
0030         if not self._use_color:
0031             return text
0032         return code + text + _RESET
0033
0034     def red(self, text: str) -> str:
0035         return self._color('\033[1;31m', text)
0036
0037     def yellow(self, text: str) -> str:
0038         return self._color('\033[1;33m', text)
0039
0040     def green(self, text: str) -> str:
0041         return self._color('\033[1;32m', text)
0042
0043     def color_len(self) -> int:
0044         """Returns the length of the color escape codes."""
0045         return len(self.red(''))
0046
0047 # Provides a default instance that prints to stdout
0048 stdout = Printer(sys.stdout)