terminal.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. # Copyright (c) 2011 The Chromium OS Authors.
  2. #
  3. # See file CREDITS for list of people who contributed to this
  4. # project.
  5. #
  6. # This program is free software; you can redistribute it and/or
  7. # modify it under the terms of the GNU General Public License as
  8. # published by the Free Software Foundation; either version 2 of
  9. # the License, or (at your option) any later version.
  10. #
  11. # This program is distributed in the hope that it will be useful,
  12. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  13. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  14. # GNU General Public License for more details.
  15. #
  16. # You should have received a copy of the GNU General Public License
  17. # along with this program; if not, write to the Free Software
  18. # Foundation, Inc., 59 Temple Place, Suite 330, Boston,
  19. # MA 02111-1307 USA
  20. #
  21. """Terminal utilities
  22. This module handles terminal interaction including ANSI color codes.
  23. """
  24. import os
  25. import sys
  26. # Selection of when we want our output to be colored
  27. COLOR_IF_TERMINAL, COLOR_ALWAYS, COLOR_NEVER = range(3)
  28. class Color(object):
  29. """Conditionally wraps text in ANSI color escape sequences."""
  30. BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)
  31. BOLD = -1
  32. COLOR_START = '\033[1;%dm'
  33. BOLD_START = '\033[1m'
  34. RESET = '\033[0m'
  35. def __init__(self, colored=COLOR_IF_TERMINAL):
  36. """Create a new Color object, optionally disabling color output.
  37. Args:
  38. enabled: True if color output should be enabled. If False then this
  39. class will not add color codes at all.
  40. """
  41. self._enabled = (colored == COLOR_ALWAYS or
  42. (colored == COLOR_IF_TERMINAL and os.isatty(sys.stdout.fileno())))
  43. def Start(self, color):
  44. """Returns a start color code.
  45. Args:
  46. color: Color to use, .e.g BLACK, RED, etc.
  47. Returns:
  48. If color is enabled, returns an ANSI sequence to start the given color,
  49. otherwise returns empty string
  50. """
  51. if self._enabled:
  52. return self.COLOR_START % (color + 30)
  53. return ''
  54. def Stop(self):
  55. """Retruns a stop color code.
  56. Returns:
  57. If color is enabled, returns an ANSI color reset sequence, otherwise
  58. returns empty string
  59. """
  60. if self._enabled:
  61. return self.RESET
  62. return ''
  63. def Color(self, color, text):
  64. """Returns text with conditionally added color escape sequences.
  65. Keyword arguments:
  66. color: Text color -- one of the color constants defined in this class.
  67. text: The text to color.
  68. Returns:
  69. If self._enabled is False, returns the original text. If it's True,
  70. returns text with color escape sequences based on the value of color.
  71. """
  72. if not self._enabled:
  73. return text
  74. if color == self.BOLD:
  75. start = self.BOLD_START
  76. else:
  77. start = self.COLOR_START % (color + 30)
  78. return start + text + self.RESET