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