codekingpro/portable-devtools

1# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.2try:3    from msvcrt import get_osfhandle4except ImportError:5    def get_osfhandle(_):6        raise OSError("This isn't windows!")7 8 9from . import win3210 11# from wincon.h12class WinColor(object):13    BLACK   = 014    BLUE    = 115    GREEN   = 216    CYAN    = 317    RED     = 418    MAGENTA = 519    YELLOW  = 620    GREY    = 721 22# from wincon.h23class WinStyle(object):24    NORMAL              = 0x00 # dim text, dim background25    BRIGHT              = 0x08 # bright text, dim background26    BRIGHT_BACKGROUND   = 0x80 # dim text, bright background27 28class WinTerm(object):29 30    def __init__(self):31        self._default = win32.GetConsoleScreenBufferInfo(win32.STDOUT).wAttributes32        self.set_attrs(self._default)33        self._default_fore = self._fore34        self._default_back = self._back35        self._default_style = self._style36        # In order to emulate LIGHT_EX in windows, we borrow the BRIGHT style.37        # So that LIGHT_EX colors and BRIGHT style do not clobber each other,38        # we track them separately, since LIGHT_EX is overwritten by Fore/Back39        # and BRIGHT is overwritten by Style codes.40        self._light = 041 42    def get_attrs(self):43        return self._fore + self._back * 16 + (self._style | self._light)44 45    def set_attrs(self, value):46        self._fore = value & 747        self._back = (value >> 4) & 748        self._style = value & (WinStyle.BRIGHT | WinStyle.BRIGHT_BACKGROUND)49 50    def reset_all(self, on_stderr=None):51        self.set_attrs(self._default)52        self.set_console(attrs=self._default)53        self._light = 054 55    def fore(self, fore=None, light=False, on_stderr=False):56        if fore is None:57            fore = self._default_fore58        self._fore = fore59        # Emulate LIGHT_EX with BRIGHT Style60        if light:61            self._light |= WinStyle.BRIGHT62        else:63            self._light &= ~WinStyle.BRIGHT64        self.set_console(on_stderr=on_stderr)65 66    def back(self, back=None, light=False, on_stderr=False):67        if back is None:68            back = self._default_back69        self._back = back70        # Emulate LIGHT_EX with BRIGHT_BACKGROUND Style71        if light:72            self._light |= WinStyle.BRIGHT_BACKGROUND73        else:74            self._light &= ~WinStyle.BRIGHT_BACKGROUND75        self.set_console(on_stderr=on_stderr)76 77    def style(self, style=None, on_stderr=False):78        if style is None:79            style = self._default_style80        self._style = style81        self.set_console(on_stderr=on_stderr)82 83    def set_console(self, attrs=None, on_stderr=False):84        if attrs is None:85            attrs = self.get_attrs()86        handle = win32.STDOUT87        if on_stderr:88            handle = win32.STDERR89        win32.SetConsoleTextAttribute(handle, attrs)90 91    def get_position(self, handle):92        position = win32.GetConsoleScreenBufferInfo(handle).dwCursorPosition93        # Because Windows coordinates are 0-based,94        # and win32.SetConsoleCursorPosition expects 1-based.95        position.X += 196        position.Y += 197        return position98 99    def set_cursor_position(self, position=None, on_stderr=False):100        if position is None:101            # I'm not currently tracking the position, so there is no default.102            # position = self.get_position()103            return104        handle = win32.STDOUT105        if on_stderr:106            handle = win32.STDERR107        win32.SetConsoleCursorPosition(handle, position)108 109    def cursor_adjust(self, x, y, on_stderr=False):110        handle = win32.STDOUT111        if on_stderr:112            handle = win32.STDERR113        position = self.get_position(handle)114        adjusted_position = (position.Y + y, position.X + x)115        win32.SetConsoleCursorPosition(handle, adjusted_position, adjust=False)116 117    def erase_screen(self, mode=0, on_stderr=False):118        # 0 should clear from the cursor to the end of the screen.119        # 1 should clear from the cursor to the beginning of the screen.120        # 2 should clear the entire screen, and move cursor to (1,1)121        handle = win32.STDOUT122        if on_stderr:123            handle = win32.STDERR124        csbi = win32.GetConsoleScreenBufferInfo(handle)125        # get the number of character cells in the current buffer126        cells_in_screen = csbi.dwSize.X * csbi.dwSize.Y127        # get number of character cells before current cursor position128        cells_before_cursor = csbi.dwSize.X * csbi.dwCursorPosition.Y + csbi.dwCursorPosition.X129        if mode == 0:130            from_coord = csbi.dwCursorPosition131            cells_to_erase = cells_in_screen - cells_before_cursor132        elif mode == 1:133            from_coord = win32.COORD(0, 0)134            cells_to_erase = cells_before_cursor135        elif mode == 2:136            from_coord = win32.COORD(0, 0)137            cells_to_erase = cells_in_screen138        else:139            # invalid mode140            return141        # fill the entire screen with blanks142        win32.FillConsoleOutputCharacter(handle, ' ', cells_to_erase, from_coord)143        # now set the buffer's attributes accordingly144        win32.FillConsoleOutputAttribute(handle, self.get_attrs(), cells_to_erase, from_coord)145        if mode == 2:146            # put the cursor where needed147            win32.SetConsoleCursorPosition(handle, (1, 1))148 149    def erase_line(self, mode=0, on_stderr=False):150        # 0 should clear from the cursor to the end of the line.151        # 1 should clear from the cursor to the beginning of the line.152        # 2 should clear the entire line.153        handle = win32.STDOUT154        if on_stderr:155            handle = win32.STDERR156        csbi = win32.GetConsoleScreenBufferInfo(handle)157        if mode == 0:158            from_coord = csbi.dwCursorPosition159            cells_to_erase = csbi.dwSize.X - csbi.dwCursorPosition.X160        elif mode == 1:161            from_coord = win32.COORD(0, csbi.dwCursorPosition.Y)162            cells_to_erase = csbi.dwCursorPosition.X163        elif mode == 2:164            from_coord = win32.COORD(0, csbi.dwCursorPosition.Y)165            cells_to_erase = csbi.dwSize.X166        else:167            # invalid mode168            return169        # fill the entire screen with blanks170        win32.FillConsoleOutputCharacter(handle, ' ', cells_to_erase, from_coord)171        # now set the buffer's attributes accordingly172        win32.FillConsoleOutputAttribute(handle, self.get_attrs(), cells_to_erase, from_coord)173 174    def set_title(self, title):175        win32.SetConsoleTitle(title)176 177 178def enable_vt_processing(fd):179    if win32.windll is None or not win32.winapi_test():180        return False181 182    try:183        handle = get_osfhandle(fd)184        mode = win32.GetConsoleMode(handle)185        win32.SetConsoleMode(186            handle,187            mode | win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING,188        )189 190        mode = win32.GetConsoleMode(handle)191        if mode & win32.ENABLE_VIRTUAL_TERMINAL_PROCESSING:192            return True193    # Can get TypeError in testsuite where 'fd' is a Mock()194    except (OSError, TypeError):195        return False196