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