terminal.py 4.4 KB

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