/sqlmaparch/thirdparty/colorama/ansitowin32.py

https://bitbucket.org/acaceres/mr-injector · Python · 189 lines · 126 code · 39 blank · 24 comment · 31 complexity · f928f5f22144adec0dc505147a226bb5 MD5 · raw file

  1. import re
  2. import sys
  3. from .ansi import AnsiFore, AnsiBack, AnsiStyle, Style
  4. from .winterm import WinTerm, WinColor, WinStyle
  5. from .win32 import windll
  6. if windll is not None:
  7. winterm = WinTerm()
  8. def is_a_tty(stream):
  9. return hasattr(stream, 'isatty') and stream.isatty()
  10. class StreamWrapper(object):
  11. '''
  12. Wraps a stream (such as stdout), acting as a transparent proxy for all
  13. attribute access apart from method 'write()', which is delegated to our
  14. Converter instance.
  15. '''
  16. def __init__(self, wrapped, converter):
  17. # double-underscore everything to prevent clashes with names of
  18. # attributes on the wrapped stream object.
  19. self.__wrapped = wrapped
  20. self.__convertor = converter
  21. def __getattr__(self, name):
  22. return getattr(self.__wrapped, name)
  23. def write(self, text):
  24. self.__convertor.write(text)
  25. class AnsiToWin32(object):
  26. '''
  27. Implements a 'write()' method which, on Windows, will strip ANSI character
  28. sequences from the text, and if outputting to a tty, will convert them into
  29. win32 function calls.
  30. '''
  31. ANSI_RE = re.compile('\033\[((?:\d|;)*)([a-zA-Z])')
  32. def __init__(self, wrapped, convert=None, strip=None, autoreset=False):
  33. # The wrapped stream (normally sys.stdout or sys.stderr)
  34. self.wrapped = wrapped
  35. # should we reset colors to defaults after every .write()
  36. self.autoreset = autoreset
  37. # create the proxy wrapping our output stream
  38. self.stream = StreamWrapper(wrapped, self)
  39. on_windows = sys.platform.startswith('win')
  40. # should we strip ANSI sequences from our output?
  41. if strip is None:
  42. strip = on_windows
  43. self.strip = strip
  44. # should we should convert ANSI sequences into win32 calls?
  45. if convert is None:
  46. convert = on_windows and is_a_tty(wrapped)
  47. self.convert = convert
  48. # dict of ansi codes to win32 functions and parameters
  49. self.win32_calls = self.get_win32_calls()
  50. # are we wrapping stderr?
  51. self.on_stderr = self.wrapped is sys.stderr
  52. def should_wrap(self):
  53. '''
  54. True if this class is actually needed. If false, then the output
  55. stream will not be affected, nor will win32 calls be issued, so
  56. wrapping stdout is not actually required. This will generally be
  57. False on non-Windows platforms, unless optional functionality like
  58. autoreset has been requested using kwargs to init()
  59. '''
  60. return self.convert or self.strip or self.autoreset
  61. def get_win32_calls(self):
  62. if self.convert and winterm:
  63. return {
  64. AnsiStyle.RESET_ALL: (winterm.reset_all, ),
  65. AnsiStyle.BRIGHT: (winterm.style, WinStyle.BRIGHT),
  66. AnsiStyle.DIM: (winterm.style, WinStyle.NORMAL),
  67. AnsiStyle.NORMAL: (winterm.style, WinStyle.NORMAL),
  68. AnsiFore.BLACK: (winterm.fore, WinColor.BLACK),
  69. AnsiFore.RED: (winterm.fore, WinColor.RED),
  70. AnsiFore.GREEN: (winterm.fore, WinColor.GREEN),
  71. AnsiFore.YELLOW: (winterm.fore, WinColor.YELLOW),
  72. AnsiFore.BLUE: (winterm.fore, WinColor.BLUE),
  73. AnsiFore.MAGENTA: (winterm.fore, WinColor.MAGENTA),
  74. AnsiFore.CYAN: (winterm.fore, WinColor.CYAN),
  75. AnsiFore.WHITE: (winterm.fore, WinColor.GREY),
  76. AnsiFore.RESET: (winterm.fore, ),
  77. AnsiBack.BLACK: (winterm.back, WinColor.BLACK),
  78. AnsiBack.RED: (winterm.back, WinColor.RED),
  79. AnsiBack.GREEN: (winterm.back, WinColor.GREEN),
  80. AnsiBack.YELLOW: (winterm.back, WinColor.YELLOW),
  81. AnsiBack.BLUE: (winterm.back, WinColor.BLUE),
  82. AnsiBack.MAGENTA: (winterm.back, WinColor.MAGENTA),
  83. AnsiBack.CYAN: (winterm.back, WinColor.CYAN),
  84. AnsiBack.WHITE: (winterm.back, WinColor.GREY),
  85. AnsiBack.RESET: (winterm.back, ),
  86. }
  87. def write(self, text):
  88. if self.strip or self.convert:
  89. self.write_and_convert(text)
  90. else:
  91. self.wrapped.write(text)
  92. self.wrapped.flush()
  93. if self.autoreset:
  94. self.reset_all()
  95. def reset_all(self):
  96. if self.convert:
  97. self.call_win32('m', (0,))
  98. elif is_a_tty(self.wrapped):
  99. self.wrapped.write(Style.RESET_ALL)
  100. def write_and_convert(self, text):
  101. '''
  102. Write the given text to our wrapped stream, stripping any ANSI
  103. sequences from the text, and optionally converting them into win32
  104. calls.
  105. '''
  106. cursor = 0
  107. for match in self.ANSI_RE.finditer(text):
  108. start, end = match.span()
  109. self.write_plain_text(text, cursor, start)
  110. self.convert_ansi(*match.groups())
  111. cursor = end
  112. self.write_plain_text(text, cursor, len(text))
  113. def write_plain_text(self, text, start, end):
  114. if start < end:
  115. self.wrapped.write(text[start:end])
  116. self.wrapped.flush()
  117. def convert_ansi(self, paramstring, command):
  118. if self.convert:
  119. params = self.extract_params(paramstring)
  120. self.call_win32(command, params)
  121. def extract_params(self, paramstring):
  122. def split(paramstring):
  123. for p in paramstring.split(';'):
  124. if p != '':
  125. yield int(p)
  126. return tuple(split(paramstring))
  127. def call_win32(self, command, params):
  128. if params == []:
  129. params = [0]
  130. if command == 'm':
  131. for param in params:
  132. if param in self.win32_calls:
  133. func_args = self.win32_calls[param]
  134. func = func_args[0]
  135. args = func_args[1:]
  136. kwargs = dict(on_stderr=self.on_stderr)
  137. func(*args, **kwargs)
  138. elif command in ('H', 'f'): # set cursor position
  139. func = winterm.set_cursor_position
  140. func(params, on_stderr=self.on_stderr)
  141. elif command in ('J'):
  142. func = winterm.erase_data
  143. func(params, on_stderr=self.on_stderr)
  144. elif command == 'A':
  145. if params == () or params == None:
  146. num_rows = 1
  147. else:
  148. num_rows = params[0]
  149. func = winterm.cursor_up
  150. func(num_rows, on_stderr=self.on_stderr)