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