tools/testing/kunit/kunit_printer.py

   1 #!/usr/bin/env python3
   2 # SPDX-License-Identifier: GPL-2.0
   3 #
   4 # Utilities for printing and coloring output.
   5 #
   6 # Copyright (C) 2022, Google LLC.
   7 # Author: Daniel Latypov <[email protected]>
   8
   9 import datetime
  10 import sys
  11 import typing
  12
  13 _RESET = '\033[0;0m'
  14
  15 class Printer:
  16         """Wraps a file object, providing utilities for coloring output, etc."""
  17
  18         def __init__(self, print: bool=True, output: typing.IO[str]=sys.stdout):
  19                 self._output = output
  20                 self._print = print
  21                 if print:
  22                         self._use_color = output.isatty()
  23                 else:
  24                         self._use_color = False
  25
  26         def print(self, message: str) -> None:
  27                 if self._print:
  28                         print(message, file=self._output)
  29
  30         def print_with_timestamp(self, message: str) -> None:
  31                 ts = datetime.datetime.now().strftime('%H:%M:%S')
  32                 self.print(f'[{ts}] {message}')
  33
  34         def _color(self, code: str, text: str) -> str:
  35                 if not self._use_color:
  36                         return text
  37                 return code + text + _RESET
  38
  39         def red(self, text: str) -> str:
  40                 return self._color('\033[1;31m', text)
  41
  42         def yellow(self, text: str) -> str:
  43                 return self._color('\033[1;33m', text)
  44
  45         def green(self, text: str) -> str:
  46                 return self._color('\033[1;32m', text)
  47
  48         def color_len(self) -> int:
  49                 """Returns the length of the color escape codes."""
  50                 return len(self.red(''))
  51
  52 # Provides a default instance that prints to stdout
  53 stdout = Printer()
  54 null_printer = Printer(print=False)