You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

ansitowin32.py 9.4KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236
  1. # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
  2. import re
  3. import sys
  4. import os
  5. from .ansi import AnsiFore, AnsiBack, AnsiStyle, Style
  6. from .winterm import WinTerm, WinColor, WinStyle
  7. from .win32 import windll, winapi_test
  8. winterm = None
  9. if windll is not None:
  10. winterm = WinTerm()
  11. def is_stream_closed(stream):
  12. return not hasattr(stream, 'closed') or stream.closed
  13. def is_a_tty(stream):
  14. return hasattr(stream, 'isatty') and stream.isatty()
  15. class StreamWrapper(object):
  16. '''
  17. Wraps a stream (such as stdout), acting as a transparent proxy for all
  18. attribute access apart from method 'write()', which is delegated to our
  19. Converter instance.
  20. '''
  21. def __init__(self, wrapped, converter):
  22. # double-underscore everything to prevent clashes with names of
  23. # attributes on the wrapped stream object.
  24. self.__wrapped = wrapped
  25. self.__convertor = converter
  26. def __getattr__(self, name):
  27. return getattr(self.__wrapped, name)
  28. def write(self, text):
  29. self.__convertor.write(text)
  30. class AnsiToWin32(object):
  31. '''
  32. Implements a 'write()' method which, on Windows, will strip ANSI character
  33. sequences from the text, and if outputting to a tty, will convert them into
  34. win32 function calls.
  35. '''
  36. ANSI_CSI_RE = re.compile('\001?\033\\[((?:\\d|;)*)([a-zA-Z])\002?') # Control Sequence Introducer
  37. ANSI_OSC_RE = re.compile('\001?\033\\]((?:.|;)*?)(\x07)\002?') # Operating System Command
  38. def __init__(self, wrapped, convert=None, strip=None, autoreset=False):
  39. # The wrapped stream (normally sys.stdout or sys.stderr)
  40. self.wrapped = wrapped
  41. # should we reset colors to defaults after every .write()
  42. self.autoreset = autoreset
  43. # create the proxy wrapping our output stream
  44. self.stream = StreamWrapper(wrapped, self)
  45. on_windows = os.name == 'nt'
  46. # We test if the WinAPI works, because even if we are on Windows
  47. # we may be using a terminal that doesn't support the WinAPI
  48. # (e.g. Cygwin Terminal). In this case it's up to the terminal
  49. # to support the ANSI codes.
  50. conversion_supported = on_windows and winapi_test()
  51. # should we strip ANSI sequences from our output?
  52. if strip is None:
  53. strip = conversion_supported or (not is_stream_closed(wrapped) and not is_a_tty(wrapped))
  54. self.strip = strip
  55. # should we should convert ANSI sequences into win32 calls?
  56. if convert is None:
  57. convert = conversion_supported and not is_stream_closed(wrapped) and is_a_tty(wrapped)
  58. self.convert = convert
  59. # dict of ansi codes to win32 functions and parameters
  60. self.win32_calls = self.get_win32_calls()
  61. # are we wrapping stderr?
  62. self.on_stderr = self.wrapped is sys.stderr
  63. def should_wrap(self):
  64. '''
  65. True if this class is actually needed. If false, then the output
  66. stream will not be affected, nor will win32 calls be issued, so
  67. wrapping stdout is not actually required. This will generally be
  68. False on non-Windows platforms, unless optional functionality like
  69. autoreset has been requested using kwargs to init()
  70. '''
  71. return self.convert or self.strip or self.autoreset
  72. def get_win32_calls(self):
  73. if self.convert and winterm:
  74. return {
  75. AnsiStyle.RESET_ALL: (winterm.reset_all, ),
  76. AnsiStyle.BRIGHT: (winterm.style, WinStyle.BRIGHT),
  77. AnsiStyle.DIM: (winterm.style, WinStyle.NORMAL),
  78. AnsiStyle.NORMAL: (winterm.style, WinStyle.NORMAL),
  79. AnsiFore.BLACK: (winterm.fore, WinColor.BLACK),
  80. AnsiFore.RED: (winterm.fore, WinColor.RED),
  81. AnsiFore.GREEN: (winterm.fore, WinColor.GREEN),
  82. AnsiFore.YELLOW: (winterm.fore, WinColor.YELLOW),
  83. AnsiFore.BLUE: (winterm.fore, WinColor.BLUE),
  84. AnsiFore.MAGENTA: (winterm.fore, WinColor.MAGENTA),
  85. AnsiFore.CYAN: (winterm.fore, WinColor.CYAN),
  86. AnsiFore.WHITE: (winterm.fore, WinColor.GREY),
  87. AnsiFore.RESET: (winterm.fore, ),
  88. AnsiFore.LIGHTBLACK_EX: (winterm.fore, WinColor.BLACK, True),
  89. AnsiFore.LIGHTRED_EX: (winterm.fore, WinColor.RED, True),
  90. AnsiFore.LIGHTGREEN_EX: (winterm.fore, WinColor.GREEN, True),
  91. AnsiFore.LIGHTYELLOW_EX: (winterm.fore, WinColor.YELLOW, True),
  92. AnsiFore.LIGHTBLUE_EX: (winterm.fore, WinColor.BLUE, True),
  93. AnsiFore.LIGHTMAGENTA_EX: (winterm.fore, WinColor.MAGENTA, True),
  94. AnsiFore.LIGHTCYAN_EX: (winterm.fore, WinColor.CYAN, True),
  95. AnsiFore.LIGHTWHITE_EX: (winterm.fore, WinColor.GREY, True),
  96. AnsiBack.BLACK: (winterm.back, WinColor.BLACK),
  97. AnsiBack.RED: (winterm.back, WinColor.RED),
  98. AnsiBack.GREEN: (winterm.back, WinColor.GREEN),
  99. AnsiBack.YELLOW: (winterm.back, WinColor.YELLOW),
  100. AnsiBack.BLUE: (winterm.back, WinColor.BLUE),
  101. AnsiBack.MAGENTA: (winterm.back, WinColor.MAGENTA),
  102. AnsiBack.CYAN: (winterm.back, WinColor.CYAN),
  103. AnsiBack.WHITE: (winterm.back, WinColor.GREY),
  104. AnsiBack.RESET: (winterm.back, ),
  105. AnsiBack.LIGHTBLACK_EX: (winterm.back, WinColor.BLACK, True),
  106. AnsiBack.LIGHTRED_EX: (winterm.back, WinColor.RED, True),
  107. AnsiBack.LIGHTGREEN_EX: (winterm.back, WinColor.GREEN, True),
  108. AnsiBack.LIGHTYELLOW_EX: (winterm.back, WinColor.YELLOW, True),
  109. AnsiBack.LIGHTBLUE_EX: (winterm.back, WinColor.BLUE, True),
  110. AnsiBack.LIGHTMAGENTA_EX: (winterm.back, WinColor.MAGENTA, True),
  111. AnsiBack.LIGHTCYAN_EX: (winterm.back, WinColor.CYAN, True),
  112. AnsiBack.LIGHTWHITE_EX: (winterm.back, WinColor.GREY, True),
  113. }
  114. return dict()
  115. def write(self, text):
  116. if self.strip or self.convert:
  117. self.write_and_convert(text)
  118. else:
  119. self.wrapped.write(text)
  120. self.wrapped.flush()
  121. if self.autoreset:
  122. self.reset_all()
  123. def reset_all(self):
  124. if self.convert:
  125. self.call_win32('m', (0,))
  126. elif not self.strip and not is_stream_closed(self.wrapped):
  127. self.wrapped.write(Style.RESET_ALL)
  128. def write_and_convert(self, text):
  129. '''
  130. Write the given text to our wrapped stream, stripping any ANSI
  131. sequences from the text, and optionally converting them into win32
  132. calls.
  133. '''
  134. cursor = 0
  135. text = self.convert_osc(text)
  136. for match in self.ANSI_CSI_RE.finditer(text):
  137. start, end = match.span()
  138. self.write_plain_text(text, cursor, start)
  139. self.convert_ansi(*match.groups())
  140. cursor = end
  141. self.write_plain_text(text, cursor, len(text))
  142. def write_plain_text(self, text, start, end):
  143. if start < end:
  144. self.wrapped.write(text[start:end])
  145. self.wrapped.flush()
  146. def convert_ansi(self, paramstring, command):
  147. if self.convert:
  148. params = self.extract_params(command, paramstring)
  149. self.call_win32(command, params)
  150. def extract_params(self, command, paramstring):
  151. if command in 'Hf':
  152. params = tuple(int(p) if len(p) != 0 else 1 for p in paramstring.split(';'))
  153. while len(params) < 2:
  154. # defaults:
  155. params = params + (1,)
  156. else:
  157. params = tuple(int(p) for p in paramstring.split(';') if len(p) != 0)
  158. if len(params) == 0:
  159. # defaults:
  160. if command in 'JKm':
  161. params = (0,)
  162. elif command in 'ABCD':
  163. params = (1,)
  164. return params
  165. def call_win32(self, command, params):
  166. if command == 'm':
  167. for param in params:
  168. if param in self.win32_calls:
  169. func_args = self.win32_calls[param]
  170. func = func_args[0]
  171. args = func_args[1:]
  172. kwargs = dict(on_stderr=self.on_stderr)
  173. func(*args, **kwargs)
  174. elif command in 'J':
  175. winterm.erase_screen(params[0], on_stderr=self.on_stderr)
  176. elif command in 'K':
  177. winterm.erase_line(params[0], on_stderr=self.on_stderr)
  178. elif command in 'Hf': # cursor position - absolute
  179. winterm.set_cursor_position(params, on_stderr=self.on_stderr)
  180. elif command in 'ABCD': # cursor position - relative
  181. n = params[0]
  182. # A - up, B - down, C - forward, D - back
  183. x, y = {'A': (0, -n), 'B': (0, n), 'C': (n, 0), 'D': (-n, 0)}[command]
  184. winterm.cursor_adjust(x, y, on_stderr=self.on_stderr)
  185. def convert_osc(self, text):
  186. for match in self.ANSI_OSC_RE.finditer(text):
  187. start, end = match.span()
  188. text = text[:start] + text[end:]
  189. paramstring, command = match.groups()
  190. if command in '\x07': # \x07 = BEL
  191. params = paramstring.split(";")
  192. # 0 - change title and icon (we will only change title)
  193. # 1 - change icon (we don't support this)
  194. # 2 - change title
  195. if params[0] in '02':
  196. winterm.set_title(params[1])
  197. return text