codekingpro/portable-devtools

1# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.2import re3import sys4import os5 6from .ansi import AnsiFore, AnsiBack, AnsiStyle, Style, BEL7from .winterm import enable_vt_processing, WinTerm, WinColor, WinStyle8from .win32 import windll, winapi_test9 10 11winterm = None12if windll is not None:13    winterm = WinTerm()14 15 16class StreamWrapper(object):17    '''18    Wraps a stream (such as stdout), acting as a transparent proxy for all19    attribute access apart from method 'write()', which is delegated to our20    Converter instance.21    '''22    def __init__(self, wrapped, converter):23        # double-underscore everything to prevent clashes with names of24        # attributes on the wrapped stream object.25        self.__wrapped = wrapped26        self.__convertor = converter27 28    def __getattr__(self, name):29        return getattr(self.__wrapped, name)30 31    def __enter__(self, *args, **kwargs):32        # special method lookup bypasses __getattr__/__getattribute__, see33        # https://stackoverflow.com/questions/12632894/why-doesnt-getattr-work-with-exit34        # thus, contextlib magic methods are not proxied via __getattr__35        return self.__wrapped.__enter__(*args, **kwargs)36 37    def __exit__(self, *args, **kwargs):38        return self.__wrapped.__exit__(*args, **kwargs)39 40    def __setstate__(self, state):41        self.__dict__ = state42 43    def __getstate__(self):44        return self.__dict__45 46    def write(self, text):47        self.__convertor.write(text)48 49    def isatty(self):50        stream = self.__wrapped51        if 'PYCHARM_HOSTED' in os.environ:52            if stream is not None and (stream is sys.__stdout__ or stream is sys.__stderr__):53                return True54        try:55            stream_isatty = stream.isatty56        except AttributeError:57            return False58        else:59            return stream_isatty()60 61    @property62    def closed(self):63        stream = self.__wrapped64        try:65            return stream.closed66        # AttributeError in the case that the stream doesn't support being closed67        # ValueError for the case that the stream has already been detached when atexit runs68        except (AttributeError, ValueError):69            return True70 71 72class AnsiToWin32(object):73    '''74    Implements a 'write()' method which, on Windows, will strip ANSI character75    sequences from the text, and if outputting to a tty, will convert them into76    win32 function calls.77    '''78    ANSI_CSI_RE = re.compile('\001?\033\\[((?:\\d|;)*)([a-zA-Z])\002?')   # Control Sequence Introducer79    ANSI_OSC_RE = re.compile('\001?\033\\]([^\a]*)(\a)\002?')             # Operating System Command80 81    def __init__(self, wrapped, convert=None, strip=None, autoreset=False):82        # The wrapped stream (normally sys.stdout or sys.stderr)83        self.wrapped = wrapped84 85        # should we reset colors to defaults after every .write()86        self.autoreset = autoreset87 88        # create the proxy wrapping our output stream89        self.stream = StreamWrapper(wrapped, self)90 91        on_windows = os.name == 'nt'92        # We test if the WinAPI works, because even if we are on Windows93        # we may be using a terminal that doesn't support the WinAPI94        # (e.g. Cygwin Terminal). In this case it's up to the terminal95        # to support the ANSI codes.96        conversion_supported = on_windows and winapi_test()97        try:98            fd = wrapped.fileno()99        except Exception:100            fd = -1101        system_has_native_ansi = not on_windows or enable_vt_processing(fd)102        have_tty = not self.stream.closed and self.stream.isatty()103        need_conversion = conversion_supported and not system_has_native_ansi104 105        # should we strip ANSI sequences from our output?106        if strip is None:107            strip = need_conversion or not have_tty108        self.strip = strip109 110        # should we should convert ANSI sequences into win32 calls?111        if convert is None:112            convert = need_conversion and have_tty113        self.convert = convert114 115        # dict of ansi codes to win32 functions and parameters116        self.win32_calls = self.get_win32_calls()117 118        # are we wrapping stderr?119        self.on_stderr = self.wrapped is sys.stderr120 121    def should_wrap(self):122        '''123        True if this class is actually needed. If false, then the output124        stream will not be affected, nor will win32 calls be issued, so125        wrapping stdout is not actually required. This will generally be126        False on non-Windows platforms, unless optional functionality like127        autoreset has been requested using kwargs to init()128        '''129        return self.convert or self.strip or self.autoreset130 131    def get_win32_calls(self):132        if self.convert and winterm:133            return {134                AnsiStyle.RESET_ALL: (winterm.reset_all, ),135                AnsiStyle.BRIGHT: (winterm.style, WinStyle.BRIGHT),136                AnsiStyle.DIM: (winterm.style, WinStyle.NORMAL),137                AnsiStyle.NORMAL: (winterm.style, WinStyle.NORMAL),138                AnsiFore.BLACK: (winterm.fore, WinColor.BLACK),139                AnsiFore.RED: (winterm.fore, WinColor.RED),140                AnsiFore.GREEN: (winterm.fore, WinColor.GREEN),141                AnsiFore.YELLOW: (winterm.fore, WinColor.YELLOW),142                AnsiFore.BLUE: (winterm.fore, WinColor.BLUE),143                AnsiFore.MAGENTA: (winterm.fore, WinColor.MAGENTA),144                AnsiFore.CYAN: (winterm.fore, WinColor.CYAN),145                AnsiFore.WHITE: (winterm.fore, WinColor.GREY),146                AnsiFore.RESET: (winterm.fore, ),147                AnsiFore.LIGHTBLACK_EX: (winterm.fore, WinColor.BLACK, True),148                AnsiFore.LIGHTRED_EX: (winterm.fore, WinColor.RED, True),149                AnsiFore.LIGHTGREEN_EX: (winterm.fore, WinColor.GREEN, True),150                AnsiFore.LIGHTYELLOW_EX: (winterm.fore, WinColor.YELLOW, True),151                AnsiFore.LIGHTBLUE_EX: (winterm.fore, WinColor.BLUE, True),152                AnsiFore.LIGHTMAGENTA_EX: (winterm.fore, WinColor.MAGENTA, True),153                AnsiFore.LIGHTCYAN_EX: (winterm.fore, WinColor.CYAN, True),154                AnsiFore.LIGHTWHITE_EX: (winterm.fore, WinColor.GREY, True),155                AnsiBack.BLACK: (winterm.back, WinColor.BLACK),156                AnsiBack.RED: (winterm.back, WinColor.RED),157                AnsiBack.GREEN: (winterm.back, WinColor.GREEN),158                AnsiBack.YELLOW: (winterm.back, WinColor.YELLOW),159                AnsiBack.BLUE: (winterm.back, WinColor.BLUE),160                AnsiBack.MAGENTA: (winterm.back, WinColor.MAGENTA),161                AnsiBack.CYAN: (winterm.back, WinColor.CYAN),162                AnsiBack.WHITE: (winterm.back, WinColor.GREY),163                AnsiBack.RESET: (winterm.back, ),164                AnsiBack.LIGHTBLACK_EX: (winterm.back, WinColor.BLACK, True),165                AnsiBack.LIGHTRED_EX: (winterm.back, WinColor.RED, True),166                AnsiBack.LIGHTGREEN_EX: (winterm.back, WinColor.GREEN, True),167                AnsiBack.LIGHTYELLOW_EX: (winterm.back, WinColor.YELLOW, True),168                AnsiBack.LIGHTBLUE_EX: (winterm.back, WinColor.BLUE, True),169                AnsiBack.LIGHTMAGENTA_EX: (winterm.back, WinColor.MAGENTA, True),170                AnsiBack.LIGHTCYAN_EX: (winterm.back, WinColor.CYAN, True),171                AnsiBack.LIGHTWHITE_EX: (winterm.back, WinColor.GREY, True),172            }173        return dict()174 175    def write(self, text):176        if self.strip or self.convert:177            self.write_and_convert(text)178        else:179            self.wrapped.write(text)180            self.wrapped.flush()181        if self.autoreset:182            self.reset_all()183 184 185    def reset_all(self):186        if self.convert:187            self.call_win32('m', (0,))188        elif not self.strip and not self.stream.closed:189            self.wrapped.write(Style.RESET_ALL)190 191 192    def write_and_convert(self, text):193        '''194        Write the given text to our wrapped stream, stripping any ANSI195        sequences from the text, and optionally converting them into win32196        calls.197        '''198        cursor = 0199        text = self.convert_osc(text)200        for match in self.ANSI_CSI_RE.finditer(text):201            start, end = match.span()202            self.write_plain_text(text, cursor, start)203            self.convert_ansi(*match.groups())204            cursor = end205        self.write_plain_text(text, cursor, len(text))206 207 208    def write_plain_text(self, text, start, end):209        if start < end:210            self.wrapped.write(text[start:end])211            self.wrapped.flush()212 213 214    def convert_ansi(self, paramstring, command):215        if self.convert:216            params = self.extract_params(command, paramstring)217            self.call_win32(command, params)218 219 220    def extract_params(self, command, paramstring):221        if command in 'Hf':222            params = tuple(int(p) if len(p) != 0 else 1 for p in paramstring.split(';'))223            while len(params) < 2:224                # defaults:225                params = params + (1,)226        else:227            params = tuple(int(p) for p in paramstring.split(';') if len(p) != 0)228            if len(params) == 0:229                # defaults:230                if command in 'JKm':231                    params = (0,)232                elif command in 'ABCD':233                    params = (1,)234 235        return params236 237 238    def call_win32(self, command, params):239        if command == 'm':240            for param in params:241                if param in self.win32_calls:242                    func_args = self.win32_calls[param]243                    func = func_args[0]244                    args = func_args[1:]245                    kwargs = dict(on_stderr=self.on_stderr)246                    func(*args, **kwargs)247        elif command in 'J':248            winterm.erase_screen(params[0], on_stderr=self.on_stderr)249        elif command in 'K':250            winterm.erase_line(params[0], on_stderr=self.on_stderr)251        elif command in 'Hf':     # cursor position - absolute252            winterm.set_cursor_position(params, on_stderr=self.on_stderr)253        elif command in 'ABCD':   # cursor position - relative254            n = params[0]255            # A - up, B - down, C - forward, D - back256            x, y = {'A': (0, -n), 'B': (0, n), 'C': (n, 0), 'D': (-n, 0)}[command]257            winterm.cursor_adjust(x, y, on_stderr=self.on_stderr)258 259 260    def convert_osc(self, text):261        for match in self.ANSI_OSC_RE.finditer(text):262            start, end = match.span()263            text = text[:start] + text[end:]264            paramstring, command = match.groups()265            if command == BEL:266                if paramstring.count(";") == 1:267                    params = paramstring.split(";")268                    # 0 - change title and icon (we will only change title)269                    # 1 - change icon (we don't support this)270                    # 2 - change title271                    if params[0] in '02':272                        winterm.set_title(params[1])273        return text274 275 276    def flush(self):277        self.wrapped.flush()278