84 lines · python
1# DExTer : Debugging Experience Tester2# ~~~~~~ ~ ~~ ~ ~~3#4# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.5# See https://llvm.org/LICENSE.txt for license information.6# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception7"""Provides Windows implementation of formatted/colored console output."""8 9import sys10 11import ctypes12import ctypes.wintypes13 14from ..PrettyOutputBase import PrettyOutputBase, Stream, _lock, _null_lock15 16 17class _CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure):18 # pylint: disable=protected-access19 _fields_ = [20 ("dwSize", ctypes.wintypes._COORD),21 ("dwCursorPosition", ctypes.wintypes._COORD),22 ("wAttributes", ctypes.c_ushort),23 ("srWindow", ctypes.wintypes._SMALL_RECT),24 ("dwMaximumWindowSize", ctypes.wintypes._COORD),25 ]26 # pylint: enable=protected-access27 28 29class PrettyOutput(PrettyOutputBase):30 stdout = Stream(sys.stdout, ctypes.windll.kernel32.GetStdHandle(-11))31 stderr = Stream(sys.stderr, ctypes.windll.kernel32.GetStdHandle(-12))32 33 def __enter__(self):34 info = _CONSOLE_SCREEN_BUFFER_INFO()35 36 for s in (PrettyOutput.stdout, PrettyOutput.stderr):37 ctypes.windll.kernel32.GetConsoleScreenBufferInfo(s.os, ctypes.byref(info))38 s.orig_color = info.wAttributes39 40 return self41 42 def __exit__(self, *args):43 self._restore_orig_color(PrettyOutput.stdout)44 self._restore_orig_color(PrettyOutput.stderr)45 46 def _restore_orig_color(self, stream, lock=_lock):47 if not stream.color_enabled:48 return49 50 with lock:51 stream = self._set_valid_stream(stream)52 self.flush(stream)53 if stream.orig_color:54 ctypes.windll.kernel32.SetConsoleTextAttribute(55 stream.os, stream.orig_color56 )57 58 def _color(self, text, color, stream, lock=_lock):59 stream = self._set_valid_stream(stream)60 with lock:61 try:62 if stream.color_enabled:63 ctypes.windll.kernel32.SetConsoleTextAttribute(stream.os, color)64 self._write(text, stream)65 finally:66 if stream.color_enabled:67 self._restore_orig_color(stream, lock=_null_lock)68 69 def red_impl(self, text, stream=None, **kwargs):70 self._color(text, 12, stream, **kwargs)71 72 def yellow_impl(self, text, stream=None, **kwargs):73 self._color(text, 14, stream, **kwargs)74 75 def green_impl(self, text, stream=None, **kwargs):76 self._color(text, 10, stream, **kwargs)77 78 def blue_impl(self, text, stream=None, **kwargs):79 self._color(text, 11, stream, **kwargs)80 81 def default_impl(self, text, stream=None, **kwargs):82 stream = self._set_valid_stream(stream)83 self._color(text, stream.orig_color, stream, **kwargs)84