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