ansitowin32.py 11 KB

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