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 <dlatypov@google.com> 8 9import datetime 10import sys 11import typing 12 13_RESET = '\033[0;0m' 14 15class 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 53stdout = Printer() 54null_printer = Printer(print=False) 55