xref: /aosp_15_r20/external/autotest/utils/frozen_chromite/lib/terminal.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1*9c5db199SXin Li# -*- coding: utf-8 -*-
2*9c5db199SXin Li# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
3*9c5db199SXin Li# Use of this source code is governed by a BSD-style license that can be
4*9c5db199SXin Li# found in the LICENSE file.
5*9c5db199SXin Li
6*9c5db199SXin Li"""Terminal utilities
7*9c5db199SXin Li
8*9c5db199SXin LiThis module handles terminal interaction including ANSI color codes.
9*9c5db199SXin Li"""
10*9c5db199SXin Li
11*9c5db199SXin Lifrom __future__ import print_function
12*9c5db199SXin Li
13*9c5db199SXin Liimport os
14*9c5db199SXin Liimport sys
15*9c5db199SXin Li
16*9c5db199SXin Lifrom autotest_lib.utils.frozen_chromite.lib import cros_build_lib
17*9c5db199SXin Li
18*9c5db199SXin Li
19*9c5db199SXin Liclass Color(object):
20*9c5db199SXin Li  """Conditionally wraps text in ANSI color escape sequences."""
21*9c5db199SXin Li  BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
22*9c5db199SXin Li  BOLD = -1
23*9c5db199SXin Li  COLOR_START = '\033[1;%dm'
24*9c5db199SXin Li  BOLD_START = '\033[1m'
25*9c5db199SXin Li  RESET = '\033[0m'
26*9c5db199SXin Li
27*9c5db199SXin Li  def __init__(self, enabled=None):
28*9c5db199SXin Li    """Create a new Color object, optionally disabling color output.
29*9c5db199SXin Li
30*9c5db199SXin Li    Args:
31*9c5db199SXin Li      enabled: True if color output should be enabled. If False then this
32*9c5db199SXin Li        class will not add color codes at all.
33*9c5db199SXin Li    """
34*9c5db199SXin Li    self._enabled = enabled
35*9c5db199SXin Li    if self._enabled is None:
36*9c5db199SXin Li      self._enabled = self.UserEnabled()
37*9c5db199SXin Li      if self._enabled is None:
38*9c5db199SXin Li        self._enabled = sys.stdout.isatty()
39*9c5db199SXin Li
40*9c5db199SXin Li  def Start(self, color):
41*9c5db199SXin Li    """Returns a start color code.
42*9c5db199SXin Li
43*9c5db199SXin Li    Args:
44*9c5db199SXin Li      color: Color to use, .e.g BLACK, RED, etc.
45*9c5db199SXin Li
46*9c5db199SXin Li    Returns:
47*9c5db199SXin Li      If color is enabled, returns an ANSI sequence to start the given color,
48*9c5db199SXin Li      otherwise returns empty string
49*9c5db199SXin Li    """
50*9c5db199SXin Li    if self._enabled:
51*9c5db199SXin Li      return self.COLOR_START % (color + 30)
52*9c5db199SXin Li    return ''
53*9c5db199SXin Li
54*9c5db199SXin Li  def Stop(self):
55*9c5db199SXin Li    """Returns a stop color code.
56*9c5db199SXin Li
57*9c5db199SXin Li    Returns:
58*9c5db199SXin Li      If color is enabled, returns an ANSI color reset sequence, otherwise
59*9c5db199SXin Li      returns empty string
60*9c5db199SXin Li    """
61*9c5db199SXin Li    if self._enabled:
62*9c5db199SXin Li      return self.RESET
63*9c5db199SXin Li    return ''
64*9c5db199SXin Li
65*9c5db199SXin Li  def Color(self, color, text):
66*9c5db199SXin Li    """Returns text with conditionally added color escape sequences.
67*9c5db199SXin Li
68*9c5db199SXin Li    Keyword arguments:
69*9c5db199SXin Li      color: Text color -- one of the color constants defined in this class.
70*9c5db199SXin Li      text: The text to color.
71*9c5db199SXin Li
72*9c5db199SXin Li    Returns:
73*9c5db199SXin Li      If self._enabled is False, returns the original text. If it's True,
74*9c5db199SXin Li      returns text with color escape sequences based on the value of color.
75*9c5db199SXin Li    """
76*9c5db199SXin Li    if not self._enabled:
77*9c5db199SXin Li      return text
78*9c5db199SXin Li    if color == self.BOLD:
79*9c5db199SXin Li      start = self.BOLD_START
80*9c5db199SXin Li    else:
81*9c5db199SXin Li      start = self.COLOR_START % (color + 30)
82*9c5db199SXin Li    return start + text + self.RESET
83*9c5db199SXin Li
84*9c5db199SXin Li  @staticmethod
85*9c5db199SXin Li  def UserEnabled():
86*9c5db199SXin Li    """See if the global colorization preference is enabled ($NOCOLOR env)"""
87*9c5db199SXin Li    is_disabled = cros_build_lib.BooleanShellValue(
88*9c5db199SXin Li        os.environ.get('NOCOLOR'), msg='$NOCOLOR env var is invalid',
89*9c5db199SXin Li        default=None)
90*9c5db199SXin Li    return not is_disabled if is_disabled is not None else None
91