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