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