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)