ansiterm.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342
  1. #!/usr/bin/env python
  2. # encoding: utf-8
  3. """
  4. Emulate a vt100 terminal in cmd.exe
  5. By wrapping sys.stdout / sys.stderr with Ansiterm,
  6. the vt100 escape characters will be interpreted and
  7. the equivalent actions will be performed with Win32
  8. console commands.
  9. """
  10. import os, re, sys
  11. from waflib import Utils
  12. wlock = Utils.threading.Lock()
  13. try:
  14. from ctypes import Structure, windll, c_short, c_ushort, c_ulong, c_int, byref, c_wchar, POINTER, c_long
  15. except ImportError:
  16. class AnsiTerm(object):
  17. def __init__(self, stream):
  18. self.stream = stream
  19. try:
  20. self.errors = self.stream.errors
  21. except AttributeError:
  22. pass # python 2.5
  23. self.encoding = self.stream.encoding
  24. def write(self, txt):
  25. try:
  26. wlock.acquire()
  27. self.stream.write(txt)
  28. self.stream.flush()
  29. finally:
  30. wlock.release()
  31. def fileno(self):
  32. return self.stream.fileno()
  33. def flush(self):
  34. self.stream.flush()
  35. def isatty(self):
  36. return self.stream.isatty()
  37. else:
  38. class COORD(Structure):
  39. _fields_ = [("X", c_short), ("Y", c_short)]
  40. class SMALL_RECT(Structure):
  41. _fields_ = [("Left", c_short), ("Top", c_short), ("Right", c_short), ("Bottom", c_short)]
  42. class CONSOLE_SCREEN_BUFFER_INFO(Structure):
  43. _fields_ = [("Size", COORD), ("CursorPosition", COORD), ("Attributes", c_ushort), ("Window", SMALL_RECT), ("MaximumWindowSize", COORD)]
  44. class CONSOLE_CURSOR_INFO(Structure):
  45. _fields_ = [('dwSize', c_ulong), ('bVisible', c_int)]
  46. try:
  47. _type = unicode
  48. except NameError:
  49. _type = str
  50. to_int = lambda number, default: number and int(number) or default
  51. STD_OUTPUT_HANDLE = -11
  52. STD_ERROR_HANDLE = -12
  53. windll.kernel32.GetStdHandle.argtypes = [c_ulong]
  54. windll.kernel32.GetStdHandle.restype = c_ulong
  55. windll.kernel32.GetConsoleScreenBufferInfo.argtypes = [c_ulong, POINTER(CONSOLE_SCREEN_BUFFER_INFO)]
  56. windll.kernel32.GetConsoleScreenBufferInfo.restype = c_long
  57. windll.kernel32.SetConsoleTextAttribute.argtypes = [c_ulong, c_ushort]
  58. windll.kernel32.SetConsoleTextAttribute.restype = c_long
  59. windll.kernel32.FillConsoleOutputCharacterW.argtypes = [c_ulong, c_wchar, c_ulong, POINTER(COORD), POINTER(c_ulong)]
  60. windll.kernel32.FillConsoleOutputCharacterW.restype = c_long
  61. windll.kernel32.FillConsoleOutputAttribute.argtypes = [c_ulong, c_ushort, c_ulong, POINTER(COORD), POINTER(c_ulong) ]
  62. windll.kernel32.FillConsoleOutputAttribute.restype = c_long
  63. windll.kernel32.SetConsoleCursorPosition.argtypes = [c_ulong, POINTER(COORD) ]
  64. windll.kernel32.SetConsoleCursorPosition.restype = c_long
  65. windll.kernel32.SetConsoleCursorInfo.argtypes = [c_ulong, POINTER(CONSOLE_CURSOR_INFO)]
  66. windll.kernel32.SetConsoleCursorInfo.restype = c_long
  67. class AnsiTerm(object):
  68. """
  69. emulate a vt100 terminal in cmd.exe
  70. """
  71. def __init__(self, s):
  72. self.stream = s
  73. try:
  74. self.errors = s.errors
  75. except AttributeError:
  76. pass # python2.5
  77. self.encoding = s.encoding
  78. self.cursor_history = []
  79. handle = (s.fileno() == 2) and STD_ERROR_HANDLE or STD_OUTPUT_HANDLE
  80. self.hconsole = windll.kernel32.GetStdHandle(handle)
  81. self._sbinfo = CONSOLE_SCREEN_BUFFER_INFO()
  82. self._csinfo = CONSOLE_CURSOR_INFO()
  83. windll.kernel32.GetConsoleCursorInfo(self.hconsole, byref(self._csinfo))
  84. # just to double check that the console is usable
  85. self._orig_sbinfo = CONSOLE_SCREEN_BUFFER_INFO()
  86. r = windll.kernel32.GetConsoleScreenBufferInfo(self.hconsole, byref(self._orig_sbinfo))
  87. self._isatty = r == 1
  88. def screen_buffer_info(self):
  89. """
  90. Updates self._sbinfo and returns it
  91. """
  92. windll.kernel32.GetConsoleScreenBufferInfo(self.hconsole, byref(self._sbinfo))
  93. return self._sbinfo
  94. def clear_line(self, param):
  95. mode = param and int(param) or 0
  96. sbinfo = self.screen_buffer_info()
  97. if mode == 1: # Clear from beginning of line to cursor position
  98. line_start = COORD(0, sbinfo.CursorPosition.Y)
  99. line_length = sbinfo.Size.X
  100. elif mode == 2: # Clear entire line
  101. line_start = COORD(sbinfo.CursorPosition.X, sbinfo.CursorPosition.Y)
  102. line_length = sbinfo.Size.X - sbinfo.CursorPosition.X
  103. else: # Clear from cursor position to end of line
  104. line_start = sbinfo.CursorPosition
  105. line_length = sbinfo.Size.X - sbinfo.CursorPosition.X
  106. chars_written = c_ulong()
  107. windll.kernel32.FillConsoleOutputCharacterW(self.hconsole, c_wchar(' '), line_length, line_start, byref(chars_written))
  108. windll.kernel32.FillConsoleOutputAttribute(self.hconsole, sbinfo.Attributes, line_length, line_start, byref(chars_written))
  109. def clear_screen(self, param):
  110. mode = to_int(param, 0)
  111. sbinfo = self.screen_buffer_info()
  112. if mode == 1: # Clear from beginning of screen to cursor position
  113. clear_start = COORD(0, 0)
  114. clear_length = sbinfo.CursorPosition.X * sbinfo.CursorPosition.Y
  115. elif mode == 2: # Clear entire screen and return cursor to home
  116. clear_start = COORD(0, 0)
  117. clear_length = sbinfo.Size.X * sbinfo.Size.Y
  118. windll.kernel32.SetConsoleCursorPosition(self.hconsole, clear_start)
  119. else: # Clear from cursor position to end of screen
  120. clear_start = sbinfo.CursorPosition
  121. clear_length = ((sbinfo.Size.X - sbinfo.CursorPosition.X) + sbinfo.Size.X * (sbinfo.Size.Y - sbinfo.CursorPosition.Y))
  122. chars_written = c_ulong()
  123. windll.kernel32.FillConsoleOutputCharacterW(self.hconsole, c_wchar(' '), clear_length, clear_start, byref(chars_written))
  124. windll.kernel32.FillConsoleOutputAttribute(self.hconsole, sbinfo.Attributes, clear_length, clear_start, byref(chars_written))
  125. def push_cursor(self, param):
  126. sbinfo = self.screen_buffer_info()
  127. self.cursor_history.append(sbinfo.CursorPosition)
  128. def pop_cursor(self, param):
  129. if self.cursor_history:
  130. old_pos = self.cursor_history.pop()
  131. windll.kernel32.SetConsoleCursorPosition(self.hconsole, old_pos)
  132. def set_cursor(self, param):
  133. y, sep, x = param.partition(';')
  134. x = to_int(x, 1) - 1
  135. y = to_int(y, 1) - 1
  136. sbinfo = self.screen_buffer_info()
  137. new_pos = COORD(
  138. min(max(0, x), sbinfo.Size.X),
  139. min(max(0, y), sbinfo.Size.Y)
  140. )
  141. windll.kernel32.SetConsoleCursorPosition(self.hconsole, new_pos)
  142. def set_column(self, param):
  143. x = to_int(param, 1) - 1
  144. sbinfo = self.screen_buffer_info()
  145. new_pos = COORD(
  146. min(max(0, x), sbinfo.Size.X),
  147. sbinfo.CursorPosition.Y
  148. )
  149. windll.kernel32.SetConsoleCursorPosition(self.hconsole, new_pos)
  150. def move_cursor(self, x_offset=0, y_offset=0):
  151. sbinfo = self.screen_buffer_info()
  152. new_pos = COORD(
  153. min(max(0, sbinfo.CursorPosition.X + x_offset), sbinfo.Size.X),
  154. min(max(0, sbinfo.CursorPosition.Y + y_offset), sbinfo.Size.Y)
  155. )
  156. windll.kernel32.SetConsoleCursorPosition(self.hconsole, new_pos)
  157. def move_up(self, param):
  158. self.move_cursor(y_offset = -to_int(param, 1))
  159. def move_down(self, param):
  160. self.move_cursor(y_offset = to_int(param, 1))
  161. def move_left(self, param):
  162. self.move_cursor(x_offset = -to_int(param, 1))
  163. def move_right(self, param):
  164. self.move_cursor(x_offset = to_int(param, 1))
  165. def next_line(self, param):
  166. sbinfo = self.screen_buffer_info()
  167. self.move_cursor(
  168. x_offset = -sbinfo.CursorPosition.X,
  169. y_offset = to_int(param, 1)
  170. )
  171. def prev_line(self, param):
  172. sbinfo = self.screen_buffer_info()
  173. self.move_cursor(
  174. x_offset = -sbinfo.CursorPosition.X,
  175. y_offset = -to_int(param, 1)
  176. )
  177. def rgb2bgr(self, c):
  178. return ((c&1) << 2) | (c&2) | ((c&4)>>2)
  179. def set_color(self, param):
  180. cols = param.split(';')
  181. sbinfo = self.screen_buffer_info()
  182. attr = sbinfo.Attributes
  183. for c in cols:
  184. c = to_int(c, 0)
  185. if 29 < c < 38: # fgcolor
  186. attr = (attr & 0xfff0) | self.rgb2bgr(c - 30)
  187. elif 39 < c < 48: # bgcolor
  188. attr = (attr & 0xff0f) | (self.rgb2bgr(c - 40) << 4)
  189. elif c == 0: # reset
  190. attr = self._orig_sbinfo.Attributes
  191. elif c == 1: # strong
  192. attr |= 0x08
  193. elif c == 4: # blink not available -> bg intensity
  194. attr |= 0x80
  195. elif c == 7: # negative
  196. attr = (attr & 0xff88) | ((attr & 0x70) >> 4) | ((attr & 0x07) << 4)
  197. windll.kernel32.SetConsoleTextAttribute(self.hconsole, attr)
  198. def show_cursor(self,param):
  199. self._csinfo.bVisible = 1
  200. windll.kernel32.SetConsoleCursorInfo(self.hconsole, byref(self._csinfo))
  201. def hide_cursor(self,param):
  202. self._csinfo.bVisible = 0
  203. windll.kernel32.SetConsoleCursorInfo(self.hconsole, byref(self._csinfo))
  204. ansi_command_table = {
  205. 'A': move_up,
  206. 'B': move_down,
  207. 'C': move_right,
  208. 'D': move_left,
  209. 'E': next_line,
  210. 'F': prev_line,
  211. 'G': set_column,
  212. 'H': set_cursor,
  213. 'f': set_cursor,
  214. 'J': clear_screen,
  215. 'K': clear_line,
  216. 'h': show_cursor,
  217. 'l': hide_cursor,
  218. 'm': set_color,
  219. 's': push_cursor,
  220. 'u': pop_cursor,
  221. }
  222. # Match either the escape sequence or text not containing escape sequence
  223. ansi_tokens = re.compile('(?:\x1b\[([0-9?;]*)([a-zA-Z])|([^\x1b]+))')
  224. def write(self, text):
  225. try:
  226. wlock.acquire()
  227. if self._isatty:
  228. for param, cmd, txt in self.ansi_tokens.findall(text):
  229. if cmd:
  230. cmd_func = self.ansi_command_table.get(cmd)
  231. if cmd_func:
  232. cmd_func(self, param)
  233. else:
  234. self.writeconsole(txt)
  235. else:
  236. # no support for colors in the console, just output the text:
  237. # eclipse or msys may be able to interpret the escape sequences
  238. self.stream.write(text)
  239. finally:
  240. wlock.release()
  241. def writeconsole(self, txt):
  242. chars_written = c_ulong()
  243. writeconsole = windll.kernel32.WriteConsoleA
  244. if isinstance(txt, _type):
  245. writeconsole = windll.kernel32.WriteConsoleW
  246. # MSDN says that there is a shared buffer of 64 KB for the console
  247. # writes. Attempt to not get ERROR_NOT_ENOUGH_MEMORY, see waf issue #746
  248. done = 0
  249. todo = len(txt)
  250. chunk = 32<<10
  251. while todo != 0:
  252. doing = min(chunk, todo)
  253. buf = txt[done:done+doing]
  254. r = writeconsole(self.hconsole, buf, doing, byref(chars_written), None)
  255. if r == 0:
  256. chunk >>= 1
  257. continue
  258. done += doing
  259. todo -= doing
  260. def fileno(self):
  261. return self.stream.fileno()
  262. def flush(self):
  263. pass
  264. def isatty(self):
  265. return self._isatty
  266. if sys.stdout.isatty() or sys.stderr.isatty():
  267. handle = sys.stdout.isatty() and STD_OUTPUT_HANDLE or STD_ERROR_HANDLE
  268. console = windll.kernel32.GetStdHandle(handle)
  269. sbinfo = CONSOLE_SCREEN_BUFFER_INFO()
  270. def get_term_cols():
  271. windll.kernel32.GetConsoleScreenBufferInfo(console, byref(sbinfo))
  272. # Issue 1401 - the progress bar cannot reach the last character
  273. return sbinfo.Size.X - 1
  274. # just try and see
  275. try:
  276. import struct, fcntl, termios
  277. except ImportError:
  278. pass
  279. else:
  280. if (sys.stdout.isatty() or sys.stderr.isatty()) and os.environ.get('TERM', '') not in ('dumb', 'emacs'):
  281. FD = sys.stdout.isatty() and sys.stdout.fileno() or sys.stderr.fileno()
  282. def fun():
  283. return struct.unpack("HHHH", fcntl.ioctl(FD, termios.TIOCGWINSZ, struct.pack("HHHH", 0, 0, 0, 0)))[1]
  284. try:
  285. fun()
  286. except Exception as e:
  287. pass
  288. else:
  289. get_term_cols = fun