1"""Windows console screen buffer handlers.""" 2 3from __future__ import print_function 4 5import atexit 6import ctypes 7import re 8import sys 9 10from colorclass.codes import ANSICodeMapping, BASE_CODES 11from colorclass.core import RE_SPLIT 12 13ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004 14INVALID_HANDLE_VALUE = -1 15IS_WINDOWS = sys.platform == 'win32' 16RE_NUMBER_SEARCH = re.compile(r'\033\[([\d;]+)m') 17STD_ERROR_HANDLE = -12 18STD_OUTPUT_HANDLE = -11 19WINDOWS_CODES = { 20 '/all': -33, '/fg': -39, '/bg': -49, 21 22 'black': 0, 'red': 4, 'green': 2, 'yellow': 6, 'blue': 1, 'magenta': 5, 'cyan': 3, 'white': 7, 23 24 'bgblack': -8, 'bgred': 64, 'bggreen': 32, 'bgyellow': 96, 'bgblue': 16, 'bgmagenta': 80, 'bgcyan': 48, 25 'bgwhite': 112, 26 27 'hiblack': 8, 'hired': 12, 'higreen': 10, 'hiyellow': 14, 'hiblue': 9, 'himagenta': 13, 'hicyan': 11, 'hiwhite': 15, 28 29 'hibgblack': 128, 'hibgred': 192, 'hibggreen': 160, 'hibgyellow': 224, 'hibgblue': 144, 'hibgmagenta': 208, 30 'hibgcyan': 176, 'hibgwhite': 240, 31 32 '/black': -39, '/red': -39, '/green': -39, '/yellow': -39, '/blue': -39, '/magenta': -39, '/cyan': -39, 33 '/white': -39, '/hiblack': -39, '/hired': -39, '/higreen': -39, '/hiyellow': -39, '/hiblue': -39, '/himagenta': -39, 34 '/hicyan': -39, '/hiwhite': -39, 35 36 '/bgblack': -49, '/bgred': -49, '/bggreen': -49, '/bgyellow': -49, '/bgblue': -49, '/bgmagenta': -49, 37 '/bgcyan': -49, '/bgwhite': -49, '/hibgblack': -49, '/hibgred': -49, '/hibggreen': -49, '/hibgyellow': -49, 38 '/hibgblue': -49, '/hibgmagenta': -49, '/hibgcyan': -49, '/hibgwhite': -49, 39} 40 41 42class COORD(ctypes.Structure): 43 """COORD structure. http://msdn.microsoft.com/en-us/library/windows/desktop/ms682119.""" 44 45 _fields_ = [ 46 ('X', ctypes.c_short), 47 ('Y', ctypes.c_short), 48 ] 49 50 51class SmallRECT(ctypes.Structure): 52 """SMALL_RECT structure. http://msdn.microsoft.com/en-us/library/windows/desktop/ms686311.""" 53 54 _fields_ = [ 55 ('Left', ctypes.c_short), 56 ('Top', ctypes.c_short), 57 ('Right', ctypes.c_short), 58 ('Bottom', ctypes.c_short), 59 ] 60 61 62class ConsoleScreenBufferInfo(ctypes.Structure): 63 """CONSOLE_SCREEN_BUFFER_INFO structure. http://msdn.microsoft.com/en-us/library/windows/desktop/ms682093.""" 64 65 _fields_ = [ 66 ('dwSize', COORD), 67 ('dwCursorPosition', COORD), 68 ('wAttributes', ctypes.c_ushort), 69 ('srWindow', SmallRECT), 70 ('dwMaximumWindowSize', COORD) 71 ] 72 73 74def init_kernel32(kernel32=None): 75 """Load a unique instance of WinDLL into memory, set arg/return types, and get stdout/err handles. 76 77 1. Since we are setting DLL function argument types and return types, we need to maintain our own instance of 78 kernel32 to prevent overriding (or being overwritten by) user's own changes to ctypes.windll.kernel32. 79 2. While we're doing all this we might as well get the handles to STDOUT and STDERR streams. 80 3. If either stream has already been replaced set return value to INVALID_HANDLE_VALUE to indicate it shouldn't be 81 replaced. 82 83 :raise AttributeError: When called on a non-Windows platform. 84 85 :param kernel32: Optional mock kernel32 object. For testing. 86 87 :return: Loaded kernel32 instance, stderr handle (int), stdout handle (int). 88 :rtype: tuple 89 """ 90 if not kernel32: 91 kernel32 = ctypes.LibraryLoader(ctypes.WinDLL).kernel32 # Load our own instance. Unique memory address. 92 kernel32.GetStdHandle.argtypes = [ctypes.c_ulong] 93 kernel32.GetStdHandle.restype = ctypes.c_void_p 94 kernel32.GetConsoleScreenBufferInfo.argtypes = [ 95 ctypes.c_void_p, 96 ctypes.POINTER(ConsoleScreenBufferInfo), 97 ] 98 kernel32.GetConsoleScreenBufferInfo.restype = ctypes.c_long 99 100 # Get handles. 101 if hasattr(sys.stderr, '_original_stream'): 102 stderr = INVALID_HANDLE_VALUE 103 else: 104 stderr = kernel32.GetStdHandle(STD_ERROR_HANDLE) 105 if hasattr(sys.stdout, '_original_stream'): 106 stdout = INVALID_HANDLE_VALUE 107 else: 108 stdout = kernel32.GetStdHandle(STD_OUTPUT_HANDLE) 109 110 return kernel32, stderr, stdout 111 112 113def get_console_info(kernel32, handle): 114 """Get information about this current console window. 115 116 http://msdn.microsoft.com/en-us/library/windows/desktop/ms683231 117 https://code.google.com/p/colorama/issues/detail?id=47 118 https://bitbucket.org/pytest-dev/py/src/4617fe46/py/_io/terminalwriter.py 119 120 Windows 10 Insider since around February 2016 finally introduced support for ANSI colors. No need to replace stdout 121 and stderr streams to intercept colors and issue multiple SetConsoleTextAttribute() calls for these consoles. 122 123 :raise OSError: When GetConsoleScreenBufferInfo or GetConsoleMode API calls fail. 124 125 :param ctypes.windll.kernel32 kernel32: Loaded kernel32 instance. 126 :param int handle: stderr or stdout handle. 127 128 :return: Foreground and background colors (integers) as well as native ANSI support (bool). 129 :rtype: tuple 130 """ 131 # Query Win32 API. 132 csbi = ConsoleScreenBufferInfo() # Populated by GetConsoleScreenBufferInfo. 133 lpcsbi = ctypes.byref(csbi) 134 dword = ctypes.c_ulong() # Populated by GetConsoleMode. 135 lpdword = ctypes.byref(dword) 136 if not kernel32.GetConsoleScreenBufferInfo(handle, lpcsbi) or not kernel32.GetConsoleMode(handle, lpdword): 137 raise ctypes.WinError() 138 139 # Parse data. 140 # buffer_width = int(csbi.dwSize.X - 1) 141 # buffer_height = int(csbi.dwSize.Y) 142 # terminal_width = int(csbi.srWindow.Right - csbi.srWindow.Left) 143 # terminal_height = int(csbi.srWindow.Bottom - csbi.srWindow.Top) 144 fg_color = csbi.wAttributes % 16 145 bg_color = csbi.wAttributes & 240 146 native_ansi = bool(dword.value & ENABLE_VIRTUAL_TERMINAL_PROCESSING) 147 148 return fg_color, bg_color, native_ansi 149 150 151def bg_color_native_ansi(kernel32, stderr, stdout): 152 """Get background color and if console supports ANSI colors natively for both streams. 153 154 :param ctypes.windll.kernel32 kernel32: Loaded kernel32 instance. 155 :param int stderr: stderr handle. 156 :param int stdout: stdout handle. 157 158 :return: Background color (int) and native ANSI support (bool). 159 :rtype: tuple 160 """ 161 try: 162 if stderr == INVALID_HANDLE_VALUE: 163 raise OSError 164 bg_color, native_ansi = get_console_info(kernel32, stderr)[1:] 165 except OSError: 166 try: 167 if stdout == INVALID_HANDLE_VALUE: 168 raise OSError 169 bg_color, native_ansi = get_console_info(kernel32, stdout)[1:] 170 except OSError: 171 bg_color, native_ansi = WINDOWS_CODES['black'], False 172 return bg_color, native_ansi 173 174 175class WindowsStream(object): 176 """Replacement stream which overrides sys.stdout or sys.stderr. When writing or printing, ANSI codes are converted. 177 178 ANSI (Linux/Unix) color codes are converted into win32 system calls, changing the next character's color before 179 printing it. Resources referenced: 180 https://github.com/tartley/colorama 181 http://www.cplusplus.com/articles/2ywTURfi/ 182 http://thomasfischer.biz/python-and-windows-terminal-colors/ 183 http://stackoverflow.com/questions/17125440/c-win32-console-color 184 http://www.tysos.org/svn/trunk/mono/corlib/System/WindowsConsoleDriver.cs 185 http://stackoverflow.com/questions/287871/print-in-terminal-with-colors-using-python 186 http://msdn.microsoft.com/en-us/library/windows/desktop/ms682088#_win32_character_attributes 187 188 :cvar list ALL_BG_CODES: List of bg Windows codes. Used to determine if requested color is foreground or background. 189 :cvar dict COMPILED_CODES: Translation dict. Keys are ANSI codes (values of BASE_CODES), values are Windows codes. 190 :ivar int default_fg: Foreground Windows color code at the time of instantiation. 191 :ivar int default_bg: Background Windows color code at the time of instantiation. 192 """ 193 194 ALL_BG_CODES = [v for k, v in WINDOWS_CODES.items() if k.startswith('bg') or k.startswith('hibg')] 195 COMPILED_CODES = dict((v, WINDOWS_CODES[k]) for k, v in BASE_CODES.items() if k in WINDOWS_CODES) 196 197 def __init__(self, kernel32, stream_handle, original_stream): 198 """Constructor. 199 200 :param ctypes.windll.kernel32 kernel32: Loaded kernel32 instance. 201 :param int stream_handle: stderr or stdout handle. 202 :param original_stream: sys.stderr or sys.stdout before being overridden by this class' instance. 203 """ 204 self._kernel32 = kernel32 205 self._stream_handle = stream_handle 206 self._original_stream = original_stream 207 self.default_fg, self.default_bg = self.colors 208 209 def __getattr__(self, item): 210 """If an attribute/function/etc is not defined in this function, retrieve the one from the original stream. 211 212 Fixes ipython arrow key presses. 213 """ 214 return getattr(self._original_stream, item) 215 216 @property 217 def colors(self): 218 """Return the current foreground and background colors.""" 219 try: 220 return get_console_info(self._kernel32, self._stream_handle)[:2] 221 except OSError: 222 return WINDOWS_CODES['white'], WINDOWS_CODES['black'] 223 224 @colors.setter 225 def colors(self, color_code): 226 """Change the foreground and background colors for subsequently printed characters. 227 228 None resets colors to their original values (when class was instantiated). 229 230 Since setting a color requires including both foreground and background codes (merged), setting just the 231 foreground color resets the background color to black, and vice versa. 232 233 This function first gets the current background and foreground colors, merges in the requested color code, and 234 sets the result. 235 236 However if we need to remove just the foreground color but leave the background color the same (or vice versa) 237 such as when {/red} is used, we must merge the default foreground color with the current background color. This 238 is the reason for those negative values. 239 240 :param int color_code: Color code from WINDOWS_CODES. 241 """ 242 if color_code is None: 243 color_code = WINDOWS_CODES['/all'] 244 245 # Get current color code. 246 current_fg, current_bg = self.colors 247 248 # Handle special negative codes. Also determine the final color code. 249 if color_code == WINDOWS_CODES['/fg']: 250 final_color_code = self.default_fg | current_bg # Reset the foreground only. 251 elif color_code == WINDOWS_CODES['/bg']: 252 final_color_code = current_fg | self.default_bg # Reset the background only. 253 elif color_code == WINDOWS_CODES['/all']: 254 final_color_code = self.default_fg | self.default_bg # Reset both. 255 elif color_code == WINDOWS_CODES['bgblack']: 256 final_color_code = current_fg # Black background. 257 else: 258 new_is_bg = color_code in self.ALL_BG_CODES 259 final_color_code = color_code | (current_fg if new_is_bg else current_bg) 260 261 # Set new code. 262 self._kernel32.SetConsoleTextAttribute(self._stream_handle, final_color_code) 263 264 def write(self, p_str): 265 """Write to stream. 266 267 :param str p_str: string to print. 268 """ 269 for segment in RE_SPLIT.split(p_str): 270 if not segment: 271 # Empty string. p_str probably starts with colors so the first item is always ''. 272 continue 273 if not RE_SPLIT.match(segment): 274 # No color codes, print regular text. 275 print(segment, file=self._original_stream, end='') 276 self._original_stream.flush() 277 continue 278 for color_code in (int(c) for c in RE_NUMBER_SEARCH.findall(segment)[0].split(';')): 279 if color_code in self.COMPILED_CODES: 280 self.colors = self.COMPILED_CODES[color_code] 281 282 283class Windows(object): 284 """Enable and disable Windows support for ANSI color character codes. 285 286 Call static method Windows.enable() to enable color support for the remainder of the process' lifetime. 287 288 This class is also a context manager. You can do this: 289 with Windows(): 290 print(Color('{autored}Test{/autored}')) 291 292 Or this: 293 with Windows(auto_colors=True): 294 print(Color('{autored}Test{/autored}')) 295 """ 296 297 @classmethod 298 def disable(cls): 299 """Restore sys.stderr and sys.stdout to their original objects. Resets colors to their original values. 300 301 :return: If streams restored successfully. 302 :rtype: bool 303 """ 304 # Skip if not on Windows. 305 if not IS_WINDOWS: 306 return False 307 308 # Restore default colors. 309 if hasattr(sys.stderr, '_original_stream'): 310 getattr(sys, 'stderr').color = None 311 if hasattr(sys.stdout, '_original_stream'): 312 getattr(sys, 'stdout').color = None 313 314 # Restore original streams. 315 changed = False 316 if hasattr(sys.stderr, '_original_stream'): 317 changed = True 318 sys.stderr = getattr(sys.stderr, '_original_stream') 319 if hasattr(sys.stdout, '_original_stream'): 320 changed = True 321 sys.stdout = getattr(sys.stdout, '_original_stream') 322 323 return changed 324 325 @staticmethod 326 def is_enabled(): 327 """Return True if either stderr or stdout has colors enabled.""" 328 return hasattr(sys.stderr, '_original_stream') or hasattr(sys.stdout, '_original_stream') 329 330 @classmethod 331 def enable(cls, auto_colors=False, reset_atexit=False): 332 """Enable color text with print() or sys.stdout.write() (stderr too). 333 334 :param bool auto_colors: Automatically selects dark or light colors based on current terminal's background 335 color. Only works with {autored} and related tags. 336 :param bool reset_atexit: Resets original colors upon Python exit (in case you forget to reset it yourself with 337 a closing tag). Does nothing on native ANSI consoles. 338 339 :return: If streams replaced successfully. 340 :rtype: bool 341 """ 342 if not IS_WINDOWS: 343 return False # Windows only. 344 345 # Get values from init_kernel32(). 346 kernel32, stderr, stdout = init_kernel32() 347 if stderr == INVALID_HANDLE_VALUE and stdout == INVALID_HANDLE_VALUE: 348 return False # No valid handles, nothing to do. 349 350 # Get console info. 351 bg_color, native_ansi = bg_color_native_ansi(kernel32, stderr, stdout) 352 353 # Set auto colors: 354 if auto_colors: 355 if bg_color in (112, 96, 240, 176, 224, 208, 160): 356 ANSICodeMapping.set_light_background() 357 else: 358 ANSICodeMapping.set_dark_background() 359 360 # Don't replace streams if ANSI codes are natively supported. 361 if native_ansi: 362 return False 363 364 # Reset on exit if requested. 365 if reset_atexit: 366 atexit.register(cls.disable) 367 368 # Overwrite stream references. 369 if stderr != INVALID_HANDLE_VALUE: 370 sys.stderr.flush() 371 sys.stderr = WindowsStream(kernel32, stderr, sys.stderr) 372 if stdout != INVALID_HANDLE_VALUE: 373 sys.stdout.flush() 374 sys.stdout = WindowsStream(kernel32, stdout, sys.stdout) 375 376 return True 377 378 def __init__(self, auto_colors=False): 379 """Constructor.""" 380 self.auto_colors = auto_colors 381 382 def __enter__(self): 383 """Context manager, enables colors on Windows.""" 384 self.enable(auto_colors=self.auto_colors) 385 386 def __exit__(self, *_): 387 """Context manager, disabled colors on Windows.""" 388 self.disable() 389