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