terminal.py 4.4 KB

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