src/colorama/ansitowin32.py

   1 # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
   2 import re
   3 import sys
   4 import os
   5
   6 from .ansi import AnsiFore, AnsiBack, AnsiStyle, Style
   7 from .winterm import WinTerm, WinColor, WinStyle
   8 from .win32 import windll, winapi_test
   9
  10
  11 winterm = None
  12 if windll is not None:
  13     winterm = WinTerm()
  14
  15
  16 def is_stream_closed(stream):
  17     return not hasattr(stream, 'closed') or stream.closed
  18
  19
  20 def is_a_tty(stream):
  21     return hasattr(stream, 'isatty') and stream.isatty()
  22
  23
  24 class StreamWrapper(object):
  25     '''
  26     Wraps a stream (such as stdout), acting as a transparent proxy for all
  27     attribute access apart from method 'write()', which is delegated to our
  28     Converter instance.
  29     '''
  30     def __init__(self, wrapped, converter):
  31         # double-underscore everything to prevent clashes with names of
  32         # attributes on the wrapped stream object.
  33         self.__wrapped = wrapped
  34         self.__convertor = converter
  35
  36     def __getattr__(self, name):
  37         return getattr(self.__wrapped, name)
  38
  39     def write(self, text):
  40         self.__convertor.write(text)
  41
  42
  43 class AnsiToWin32(object):
  44     '''
  45     Implements a 'write()' method which, on Windows, will strip ANSI character
  46     sequences from the text, and if outputting to a tty, will convert them into
  47     win32 function calls.
  48     '''
  49     ANSI_CSI_RE = re.compile('\001?\033\[((?:\d|;)*)([a-zA-Z])\002?')     # Control Sequence Introducer
  50     ANSI_OSC_RE = re.compile('\001?\033\]((?:.|;)*?)(\x07)\002?')         # Operating System Command
  51
  52     def __init__(self, wrapped, convert=None, strip=None, autoreset=False):
  53         # The wrapped stream (normally sys.stdout or sys.stderr)
  54         self.wrapped = wrapped
  55
  56         # should we reset colors to defaults after every .write()
  57         self.autoreset = autoreset
  58
  59         # create the proxy wrapping our output stream
  60         self.stream = StreamWrapper(wrapped, self)
  61
  62         on_windows = os.name == 'nt'
  63         # We test if the WinAPI works, because even if we are on Windows
  64         # we may be using a terminal that doesn't support the WinAPI
  65         # (e.g. Cygwin Terminal). In this case it's up to the terminal
  66         # to support the ANSI codes.
  67         conversion_supported = on_windows and winapi_test()
  68
  69         # should we strip ANSI sequences from our output?
  70         if strip is None:
  71             strip = conversion_supported or (not is_stream_closed(wrapped) and not is_a_tty(wrapped))
  72         self.strip = strip
  73
  74         # should we should convert ANSI sequences into win32 calls?
  75         if convert is None:
  76             convert = conversion_supported and not is_stream_closed(wrapped) and is_a_tty(wrapped)
  77         self.convert = convert
  78
  79         # dict of ansi codes to win32 functions and parameters
  80         self.win32_calls = self.get_win32_calls()
  81
  82         # are we wrapping stderr?
  83         self.on_stderr = self.wrapped is sys.stderr
  84
  85     def should_wrap(self):
  86         '''
  87         True if this class is actually needed. If false, then the output
  88         stream will not be affected, nor will win32 calls be issued, so
  89         wrapping stdout is not actually required. This will generally be
  90         False on non-Windows platforms, unless optional functionality like
  91         autoreset has been requested using kwargs to init()
  92         '''
  93         return self.convert or self.strip or self.autoreset
  94
  95     def get_win32_calls(self):
  96         if self.convert and winterm:
  97             return {
  98                 AnsiStyle.RESET_ALL: (winterm.reset_all, ),
  99                 AnsiStyle.BRIGHT: (winterm.style, WinStyle.BRIGHT),
 100                 AnsiStyle.DIM: (winterm.style, WinStyle.NORMAL),
 101                 AnsiStyle.NORMAL: (winterm.style, WinStyle.NORMAL),
 102                 AnsiFore.BLACK: (winterm.fore, WinColor.BLACK),
 103                 AnsiFore.RED: (winterm.fore, WinColor.RED),
 104                 AnsiFore.GREEN: (winterm.fore, WinColor.GREEN),
 105                 AnsiFore.YELLOW: (winterm.fore, WinColor.YELLOW),
 106                 AnsiFore.BLUE: (winterm.fore, WinColor.BLUE),
 107                 AnsiFore.MAGENTA: (winterm.fore, WinColor.MAGENTA),
 108                 AnsiFore.CYAN: (winterm.fore, WinColor.CYAN),
 109                 AnsiFore.WHITE: (winterm.fore, WinColor.GREY),
 110                 AnsiFore.RESET: (winterm.fore, ),
 111                 AnsiFore.LIGHTBLACK_EX: (winterm.fore, WinColor.BLACK, True),
 112                 AnsiFore.LIGHTRED_EX: (winterm.fore, WinColor.RED, True),
 113                 AnsiFore.LIGHTGREEN_EX: (winterm.fore, WinColor.GREEN, True),
 114                 AnsiFore.LIGHTYELLOW_EX: (winterm.fore, WinColor.YELLOW, True),
 115                 AnsiFore.LIGHTBLUE_EX: (winterm.fore, WinColor.BLUE, True),
 116                 AnsiFore.LIGHTMAGENTA_EX: (winterm.fore, WinColor.MAGENTA, True),
 117                 AnsiFore.LIGHTCYAN_EX: (winterm.fore, WinColor.CYAN, True),
 118                 AnsiFore.LIGHTWHITE_EX: (winterm.fore, WinColor.GREY, True),
 119                 AnsiBack.BLACK: (winterm.back, WinColor.BLACK),
 120                 AnsiBack.RED: (winterm.back, WinColor.RED),
 121                 AnsiBack.GREEN: (winterm.back, WinColor.GREEN),
 122                 AnsiBack.YELLOW: (winterm.back, WinColor.YELLOW),
 123                 AnsiBack.BLUE: (winterm.back, WinColor.BLUE),
 124                 AnsiBack.MAGENTA: (winterm.back, WinColor.MAGENTA),
 125                 AnsiBack.CYAN: (winterm.back, WinColor.CYAN),
 126                 AnsiBack.WHITE: (winterm.back, WinColor.GREY),
 127                 AnsiBack.RESET: (winterm.back, ),
 128                 AnsiBack.LIGHTBLACK_EX: (winterm.back, WinColor.BLACK, True),
 129                 AnsiBack.LIGHTRED_EX: (winterm.back, WinColor.RED, True),
 130                 AnsiBack.LIGHTGREEN_EX: (winterm.back, WinColor.GREEN, True),
 131                 AnsiBack.LIGHTYELLOW_EX: (winterm.back, WinColor.YELLOW, True),
 132                 AnsiBack.LIGHTBLUE_EX: (winterm.back, WinColor.BLUE, True),
 133                 AnsiBack.LIGHTMAGENTA_EX: (winterm.back, WinColor.MAGENTA, True),
 134                 AnsiBack.LIGHTCYAN_EX: (winterm.back, WinColor.CYAN, True),
 135                 AnsiBack.LIGHTWHITE_EX: (winterm.back, WinColor.GREY, True),
 136             }
 137         return dict()
 138
 139     def write(self, text):
 140         if self.strip or self.convert:
 141             self.write_and_convert(text)
 142         else:
 143             self.wrapped.write(text)
 144             self.wrapped.flush()
 145         if self.autoreset:
 146             self.reset_all()
 147
 148
 149     def reset_all(self):
 150         if self.convert:
 151             self.call_win32('m', (0,))
 152         elif not self.strip and not is_stream_closed(self.wrapped):
 153             self.wrapped.write(Style.RESET_ALL)
 154
 155
 156     def write_and_convert(self, text):
 157         '''
 158         Write the given text to our wrapped stream, stripping any ANSI
 159         sequences from the text, and optionally converting them into win32
 160         calls.
 161         '''
 162         cursor = 0
 163         text = self.convert_osc(text)
 164         for match in self.ANSI_CSI_RE.finditer(text):
 165             start, end = match.span()
 166             self.write_plain_text(text, cursor, start)
 167             self.convert_ansi(*match.groups())
 168             cursor = end
 169         self.write_plain_text(text, cursor, len(text))
 170
 171
 172     def write_plain_text(self, text, start, end):
 173         if start < end:
 174             self.wrapped.write(text[start:end])
 175             self.wrapped.flush()
 176
 177
 178     def convert_ansi(self, paramstring, command):
 179         if self.convert:
 180             params = self.extract_params(command, paramstring)
 181             self.call_win32(command, params)
 182
 183
 184     def extract_params(self, command, paramstring):
 185         if command in 'Hf':
 186             params = tuple(int(p) if len(p) != 0 else 1 for p in paramstring.split(';'))
 187             while len(params) < 2:
 188                 # defaults:
 189                 params = params + (1,)
 190         else:
 191             params = tuple(int(p) for p in paramstring.split(';') if len(p) != 0)
 192             if len(params) == 0:
 193                 # defaults:
 194                 if command in 'JKm':
 195                     params = (0,)
 196                 elif command in 'ABCD':
 197                     params = (1,)
 198
 199         return params
 200
 201
 202     def call_win32(self, command, params):
 203         if command == 'm':
 204             for param in params:
 205                 if param in self.win32_calls:
 206                     func_args = self.win32_calls[param]
 207                     func = func_args[0]
 208                     args = func_args[1:]
 209                     kwargs = dict(on_stderr=self.on_stderr)
 210                     func(*args, **kwargs)
 211         elif command in 'J':
 212             winterm.erase_screen(params[0], on_stderr=self.on_stderr)
 213         elif command in 'K':
 214             winterm.erase_line(params[0], on_stderr=self.on_stderr)
 215         elif command in 'Hf':     # cursor position - absolute
 216             winterm.set_cursor_position(params, on_stderr=self.on_stderr)
 217         elif command in 'ABCD':   # cursor position - relative
 218             n = params[0]
 219             # A - up, B - down, C - forward, D - back
 220             x, y = {'A': (0, -n), 'B': (0, n), 'C': (n, 0), 'D': (-n, 0)}[command]
 221             winterm.cursor_adjust(x, y, on_stderr=self.on_stderr)
 222
 223
 224     def convert_osc(self, text):
 225         for match in self.ANSI_OSC_RE.finditer(text):
 226             start, end = match.span()
 227             text = text[:start] + text[end:]
 228             paramstring, command = match.groups()
 229             if command in '\x07':       # \x07 = BEL
 230                 params = paramstring.split(";")
 231                 # 0 - change title and icon (we will only change title)
 232                 # 1 - change icon (we don't support this)
 233                 # 2 - change title
 234                 if params[0] in '02':
 235                     winterm.set_title(params[1])
 236         return text