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