ansistrm.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. #
  2. # Copyright (C) 2010 Vinay Sajip. All rights reserved.
  3. # https://gist.github.com/758430
  4. import ctypes
  5. import logging
  6. import os
  7. class ColorizingStreamHandler(logging.StreamHandler):
  8. # color names to indices
  9. color_map = {
  10. 'black': 0,
  11. 'red': 1,
  12. 'green': 2,
  13. 'yellow': 3,
  14. 'blue': 4,
  15. 'magenta': 5,
  16. 'cyan': 6,
  17. 'white': 7,
  18. }
  19. #levels to (background, foreground, bold/intense)
  20. if os.name == 'nt':
  21. level_map = {
  22. logging.DEBUG: (None, 'blue', True),
  23. logging.INFO: (None, 'white', False),
  24. logging.INFO+5: (None, 'green', True), #log.interest
  25. logging.WARNING: (None, 'yellow', True),
  26. logging.WARNING+5: (None, 'cyan', True), #log.event
  27. logging.ERROR: (None, 'red', True),
  28. logging.CRITICAL: ('red', 'white', True),
  29. }
  30. else:
  31. level_map = {
  32. logging.DEBUG: (None, 'blue', False),
  33. logging.INFO: (None, 'white', False),
  34. logging.INFO+5: (None, 'green', True), #log.interest
  35. logging.WARNING: (None, 'yellow', False),
  36. logging.WARNING+5: (None, 'cyan', True), #log.event
  37. logging.ERROR: (None, 'red', False),
  38. logging.CRITICAL: ('red', 'white', True),
  39. }
  40. csi = '\x1b['
  41. reset = '\x1b[0m'
  42. @property
  43. def is_tty(self):
  44. isatty = getattr(self.stream, 'isatty', None)
  45. return isatty and isatty()
  46. def emit(self, record):
  47. try:
  48. message = self.format(record)
  49. if isinstance(message, unicode):
  50. message = message.encode('utf-8')
  51. stream = self.stream
  52. if not self.is_tty:
  53. stream.write(message)
  54. else:
  55. self.output_colorized(message)
  56. stream.write(getattr(self, 'terminator', '\n'))
  57. self.flush()
  58. except (KeyboardInterrupt, SystemExit):
  59. raise
  60. except:
  61. self.handleError(record)
  62. if os.name != 'nt':
  63. def output_colorized(self, message):
  64. self.stream.write(message)
  65. else:
  66. import re
  67. ansi_esc = re.compile(r'\x1b\[((?:\d+)(?:;(?:\d+))*)m')
  68. nt_color_map = {
  69. 0: 0x00, # black
  70. 1: 0x04, # red
  71. 2: 0x02, # green
  72. 3: 0x06, # yellow
  73. 4: 0x01, # blue
  74. 5: 0x05, # magenta
  75. 6: 0x03, # cyan
  76. 7: 0x07, # white
  77. }
  78. def output_colorized(self, message):
  79. parts = self.ansi_esc.split(message)
  80. write = self.stream.write
  81. h = None
  82. fd = getattr(self.stream, 'fileno', None)
  83. if fd is not None:
  84. fd = fd()
  85. if fd in (1, 2): # stdout or stderr
  86. h = ctypes.windll.kernel32.GetStdHandle(-10 - fd)
  87. while parts:
  88. text = parts.pop(0)
  89. if text:
  90. write(text)
  91. if parts:
  92. params = parts.pop(0)
  93. if h is not None:
  94. params = [int(p) for p in params.split(';')]
  95. color = 0
  96. for p in params:
  97. if 40 <= p <= 47:
  98. color |= self.nt_color_map[p - 40] << 4
  99. elif 30 <= p <= 37:
  100. color |= self.nt_color_map[p - 30]
  101. elif p == 1:
  102. color |= 0x08 # foreground intensity on
  103. elif p == 0: # reset to default color
  104. color = 0x07
  105. else:
  106. pass # error condition ignored
  107. ctypes.windll.kernel32.SetConsoleTextAttribute(h, color)
  108. def colorize(self, message, record):
  109. if record.levelno in self.level_map:
  110. bg, fg, bold = self.level_map[record.levelno]
  111. params = []
  112. if bg in self.color_map:
  113. params.append(str(self.color_map[bg] + 40))
  114. if fg in self.color_map:
  115. params.append(str(self.color_map[fg] + 30))
  116. if bold:
  117. params.append('1')
  118. if params:
  119. message = ''.join((self.csi, ';'.join(params),
  120. 'm', message, self.reset))
  121. return message
  122. def format(self, record):
  123. message = logging.StreamHandler.format(self, record)
  124. if self.is_tty:
  125. # Don't colorize any traceback
  126. parts = message.split('\n', 1)
  127. parts[0] = self.colorize(parts[0], record)
  128. message = '\n'.join(parts)
  129. return message
  130. def main():
  131. root = logging.getLogger()
  132. root.setLevel(logging.DEBUG)
  133. root.addHandler(ColorizingStreamHandler())
  134. logging.debug('DEBUG')
  135. logging.info('INFO')
  136. logging.warning('WARNING')
  137. logging.error('ERROR')
  138. logging.critical('CRITICAL')
  139. if __name__ == '__main__':
  140. main()