terminal.py 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270
  1. # SPDX-License-Identifier: GPL-2.0+
  2. # Copyright (c) 2011 The Chromium OS Authors.
  3. #
  4. """Terminal utilities
  5. This module handles terminal interaction including ANSI color codes.
  6. """
  7. import os
  8. import re
  9. import shutil
  10. import sys
  11. # Selection of when we want our output to be colored
  12. COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3)
  13. # Initially, we are set up to print to the terminal
  14. print_test_mode = False
  15. print_test_list = []
  16. # The length of the last line printed without a newline
  17. last_print_len = None
  18. # credit:
  19. # stackoverflow.com/questions/14693701/how-can-i-remove-the-ansi-escape-sequences-from-a-string-in-python
  20. ansi_escape = re.compile(r'\x1b(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
  21. class PrintLine:
  22. """A line of text output
  23. Members:
  24. text: Text line that was printed
  25. newline: True to output a newline after the text
  26. colour: Text colour to use
  27. """
  28. def __init__(self, text, colour, newline=True, bright=True):
  29. self.text = text
  30. self.newline = newline
  31. self.colour = colour
  32. self.bright = bright
  33. def __eq__(self, other):
  34. return (self.text == other.text and
  35. self.newline == other.newline and
  36. self.colour == other.colour and
  37. self.bright == other.bright)
  38. def __str__(self):
  39. return ("newline=%s, colour=%s, bright=%d, text='%s'" %
  40. (self.newline, self.colour, self.bright, self.text))
  41. def CalcAsciiLen(text):
  42. """Calculate the length of a string, ignoring any ANSI sequences
  43. When displayed on a terminal, ANSI sequences don't take any space, so we
  44. need to ignore them when calculating the length of a string.
  45. Args:
  46. text: Text to check
  47. Returns:
  48. Length of text, after skipping ANSI sequences
  49. >>> col = Color(COLOR_ALWAYS)
  50. >>> text = col.Color(Color.RED, 'abc')
  51. >>> len(text)
  52. 14
  53. >>> CalcAsciiLen(text)
  54. 3
  55. >>>
  56. >>> text += 'def'
  57. >>> CalcAsciiLen(text)
  58. 6
  59. >>> text += col.Color(Color.RED, 'abc')
  60. >>> CalcAsciiLen(text)
  61. 9
  62. """
  63. result = ansi_escape.sub('', text)
  64. return len(result)
  65. def TrimAsciiLen(text, size):
  66. """Trim a string containing ANSI sequences to the given ASCII length
  67. The string is trimmed with ANSI sequences being ignored for the length
  68. calculation.
  69. >>> col = Color(COLOR_ALWAYS)
  70. >>> text = col.Color(Color.RED, 'abc')
  71. >>> len(text)
  72. 14
  73. >>> CalcAsciiLen(TrimAsciiLen(text, 4))
  74. 3
  75. >>> CalcAsciiLen(TrimAsciiLen(text, 2))
  76. 2
  77. >>> text += 'def'
  78. >>> CalcAsciiLen(TrimAsciiLen(text, 4))
  79. 4
  80. >>> text += col.Color(Color.RED, 'ghi')
  81. >>> CalcAsciiLen(TrimAsciiLen(text, 7))
  82. 7
  83. """
  84. if CalcAsciiLen(text) < size:
  85. return text
  86. pos = 0
  87. out = ''
  88. left = size
  89. # Work through each ANSI sequence in turn
  90. for m in ansi_escape.finditer(text):
  91. # Find the text before the sequence and add it to our string, making
  92. # sure it doesn't overflow
  93. before = text[pos:m.start()]
  94. toadd = before[:left]
  95. out += toadd
  96. # Figure out how much non-ANSI space we have left
  97. left -= len(toadd)
  98. # Add the ANSI sequence and move to the position immediately after it
  99. out += m.group()
  100. pos = m.start() + len(m.group())
  101. # Deal with text after the last ANSI sequence
  102. after = text[pos:]
  103. toadd = after[:left]
  104. out += toadd
  105. return out
  106. def Print(text='', newline=True, colour=None, limit_to_line=False, bright=True):
  107. """Handle a line of output to the terminal.
  108. In test mode this is recorded in a list. Otherwise it is output to the
  109. terminal.
  110. Args:
  111. text: Text to print
  112. newline: True to add a new line at the end of the text
  113. colour: Colour to use for the text
  114. """
  115. global last_print_len
  116. if print_test_mode:
  117. print_test_list.append(PrintLine(text, colour, newline, bright))
  118. else:
  119. if colour:
  120. col = Color()
  121. text = col.Color(colour, text, bright=bright)
  122. if newline:
  123. print(text)
  124. last_print_len = None
  125. else:
  126. if limit_to_line:
  127. cols = shutil.get_terminal_size().columns
  128. text = TrimAsciiLen(text, cols)
  129. print(text, end='', flush=True)
  130. last_print_len = CalcAsciiLen(text)
  131. def PrintClear():
  132. """Clear a previously line that was printed with no newline"""
  133. global last_print_len
  134. if last_print_len:
  135. print('\r%s\r' % (' '* last_print_len), end='', flush=True)
  136. last_print_len = None
  137. def SetPrintTestMode(enable=True):
  138. """Go into test mode, where all printing is recorded"""
  139. global print_test_mode
  140. print_test_mode = enable
  141. GetPrintTestLines()
  142. def GetPrintTestLines():
  143. """Get a list of all lines output through Print()
  144. Returns:
  145. A list of PrintLine objects
  146. """
  147. global print_test_list
  148. ret = print_test_list
  149. print_test_list = []
  150. return ret
  151. def EchoPrintTestLines():
  152. """Print out the text lines collected"""
  153. for line in print_test_list:
  154. if line.colour:
  155. col = Color()
  156. print(col.Color(line.colour, line.text), end='')
  157. else:
  158. print(line.text, end='')
  159. if line.newline:
  160. print()
  161. class Color(object):
  162. """Conditionally wraps text in ANSI color escape sequences."""
  163. BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
  164. BOLD = -1
  165. BRIGHT_START = '\033[1;%dm'
  166. NORMAL_START = '\033[22;%dm'
  167. BOLD_START = '\033[1m'
  168. RESET = '\033[0m'
  169. def __init__(self, colored=COLOR_IF_TERMINAL):
  170. """Create a new Color object, optionally disabling color output.
  171. Args:
  172. enabled: True if color output should be enabled. If False then this
  173. class will not add color codes at all.
  174. """
  175. try:
  176. self._enabled = (colored == COLOR_ALWAYS or
  177. (colored == COLOR_IF_TERMINAL and
  178. os.isatty(sys.stdout.fileno())))
  179. except:
  180. self._enabled = False
  181. def Start(self, color, bright=True):
  182. """Returns a start color code.
  183. Args:
  184. color: Color to use, .e.g BLACK, RED, etc.
  185. Returns:
  186. If color is enabled, returns an ANSI sequence to start the given
  187. color, otherwise returns empty string
  188. """
  189. if self._enabled:
  190. base = self.BRIGHT_START if bright else self.NORMAL_START
  191. return base % (color + 30)
  192. return ''
  193. def Stop(self):
  194. """Returns a stop color code.
  195. Returns:
  196. If color is enabled, returns an ANSI color reset sequence,
  197. otherwise returns empty string
  198. """
  199. if self._enabled:
  200. return self.RESET
  201. return ''
  202. def Color(self, color, text, bright=True):
  203. """Returns text with conditionally added color escape sequences.
  204. Keyword arguments:
  205. color: Text color -- one of the color constants defined in this
  206. class.
  207. text: The text to color.
  208. Returns:
  209. If self._enabled is False, returns the original text. If it's True,
  210. returns text with color escape sequences based on the value of
  211. color.
  212. """
  213. if not self._enabled:
  214. return text
  215. if color == self.BOLD:
  216. start = self.BOLD_START
  217. else:
  218. base = self.BRIGHT_START if bright else self.NORMAL_START
  219. start = base % (color + 30)
  220. return start + text + self.RESET