xref: /openbmc/u-boot/tools/patman/terminal.py (revision 6e87ae1c) 
1# Copyright (c) 2011 The Chromium OS Authors.
2#
3# SPDX-License-Identifier:	GPL-2.0+
4#
5
6"""Terminal utilities
7
8This module handles terminal interaction including ANSI color codes.
9"""
10
11from __future__ import print_function
12
13import os
14import sys
15
16# Selection of when we want our output to be colored
17COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3)
18
19# Initially, we are set up to print to the terminal
20print_test_mode = False
21print_test_list = []
22
23class PrintLine:
24    """A line of text output
25
26    Members:
27        text: Text line that was printed
28        newline: True to output a newline after the text
29        colour: Text colour to use
30    """
31    def __init__(self, text, newline, colour):
32        self.text = text
33        self.newline = newline
34        self.colour = colour
35
36    def __str__(self):
37        return 'newline=%s, colour=%s, text=%s' % (self.newline, self.colour,
38                self.text)
39
40def Print(text='', newline=True, colour=None):
41    """Handle a line of output to the terminal.
42
43    In test mode this is recorded in a list. Otherwise it is output to the
44    terminal.
45
46    Args:
47        text: Text to print
48        newline: True to add a new line at the end of the text
49        colour: Colour to use for the text
50    """
51    if print_test_mode:
52        print_test_list.append(PrintLine(text, newline, colour))
53    else:
54        if colour:
55            col = Color()
56            text = col.Color(colour, text)
57        print(text, end='')
58        if newline:
59            print()
60        else:
61            sys.stdout.flush()
62
63def SetPrintTestMode():
64    """Go into test mode, where all printing is recorded"""
65    global print_test_mode
66
67    print_test_mode = True
68
69def GetPrintTestLines():
70    """Get a list of all lines output through Print()
71
72    Returns:
73        A list of PrintLine objects
74    """
75    global print_test_list
76
77    ret = print_test_list
78    print_test_list = []
79    return ret
80
81def EchoPrintTestLines():
82    """Print out the text lines collected"""
83    for line in print_test_list:
84        if line.colour:
85            col = Color()
86            print(col.Color(line.colour, line.text), end='')
87        else:
88            print(line.text, end='')
89        if line.newline:
90            print()
91
92
93class Color(object):
94    """Conditionally wraps text in ANSI color escape sequences."""
95    BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
96    BOLD = -1
97    BRIGHT_START = '\033[1;%dm'
98    NORMAL_START = '\033[22;%dm'
99    BOLD_START = '\033[1m'
100    RESET = '\033[0m'
101
102    def __init__(self, colored=COLOR_IF_TERMINAL):
103        """Create a new Color object, optionally disabling color output.
104
105        Args:
106          enabled: True if color output should be enabled. If False then this
107            class will not add color codes at all.
108        """
109        try:
110            self._enabled = (colored == COLOR_ALWAYS or
111                    (colored == COLOR_IF_TERMINAL and
112                     os.isatty(sys.stdout.fileno())))
113        except:
114            self._enabled = False
115
116    def Start(self, color, bright=True):
117        """Returns a start color code.
118
119        Args:
120          color: Color to use, .e.g BLACK, RED, etc.
121
122        Returns:
123          If color is enabled, returns an ANSI sequence to start the given
124          color, otherwise returns empty string
125        """
126        if self._enabled:
127            base = self.BRIGHT_START if bright else self.NORMAL_START
128            return base % (color + 30)
129        return ''
130
131    def Stop(self):
132        """Retruns a stop color code.
133
134        Returns:
135          If color is enabled, returns an ANSI color reset sequence,
136          otherwise returns empty string
137        """
138        if self._enabled:
139            return self.RESET
140        return ''
141
142    def Color(self, color, text, bright=True):
143        """Returns text with conditionally added color escape sequences.
144
145        Keyword arguments:
146          color: Text color -- one of the color constants defined in this
147                  class.
148          text: The text to color.
149
150        Returns:
151          If self._enabled is False, returns the original text. If it's True,
152          returns text with color escape sequences based on the value of
153          color.
154        """
155        if not self._enabled:
156            return text
157        if color == self.BOLD:
158            start = self.BOLD_START
159        else:
160            base = self.BRIGHT_START if bright else self.NORMAL_START
161            start = base % (color + 30)
162        return start + text + self.RESET
163