Skip to content

Commit

Permalink
Support colors in jupyter notebooks.
Browse files Browse the repository at this point in the history
  • Loading branch information
@raldone01
raldone01 committed Oct 21, 2024
1 parent 1368087 commit d2a35c6
Showing 1 changed file with 40 additions and 37 deletions.
77 changes: 40 additions & 37 deletions colorama/ansitowin32.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,12 @@


class StreamWrapper(object):
'''
"""
Wraps a stream (such as stdout), acting as a transparent proxy for all
attribute access apart from method 'write()', which is delegated to our
Converter instance.
'''
"""

def __init__(self, wrapped, converter):
# double-underscore everything to prevent clashes with names of
# attributes on the wrapped stream object.
Expand Down Expand Up @@ -48,9 +49,14 @@ def write(self, text):

def isatty(self):
stream = self.__wrapped
if 'PYCHARM_HOSTED' in os.environ:
if stream is not None and (stream is sys.__stdout__ or stream is sys.__stderr__):
if "PYCHARM_HOSTED" in os.environ:
if stream is not None and (
stream is sys.__stdout__ or stream is sys.__stderr__
):
return True
# Detect if we are running in a jupyter notebook which supports ANSI colors
elif "JPY_PARENT_PID" in os.environ:
return False
try:
stream_isatty = stream.isatty
except AttributeError:
Expand All @@ -70,13 +76,18 @@ def closed(self):


class AnsiToWin32(object):
'''
"""
Implements a 'write()' method which, on Windows, will strip ANSI character
sequences from the text, and if outputting to a tty, will convert them into
win32 function calls.
'''
ANSI_CSI_RE = re.compile('\001?\033\\[((?:\\d|;)*)([a-zA-Z])\002?') # Control Sequence Introducer
ANSI_OSC_RE = re.compile('\001?\033\\]([^\a]*)(\a)\002?') # Operating System Command
"""

ANSI_CSI_RE = re.compile(
"\001?\033\\[((?:\\d|;)*)([a-zA-Z])\002?"
) # Control Sequence Introducer
ANSI_OSC_RE = re.compile(
"\001?\033\\]([^\a]*)(\a)\002?"
) # Operating System Command

def __init__(self, wrapped, convert=None, strip=None, autoreset=False):
# The wrapped stream (normally sys.stdout or sys.stderr)
Expand All @@ -88,7 +99,7 @@ def __init__(self, wrapped, convert=None, strip=None, autoreset=False):
# create the proxy wrapping our output stream
self.stream = StreamWrapper(wrapped, self)

on_windows = os.name == 'nt'
on_windows = os.name == "nt"
# We test if the WinAPI works, because even if we are on Windows
# we may be using a terminal that doesn't support the WinAPI
# (e.g. Cygwin Terminal). In this case it's up to the terminal
Expand Down Expand Up @@ -119,19 +130,19 @@ def __init__(self, wrapped, convert=None, strip=None, autoreset=False):
self.on_stderr = self.wrapped is sys.stderr

def should_wrap(self):
'''
"""
True if this class is actually needed. If false, then the output
stream will not be affected, nor will win32 calls be issued, so
wrapping stdout is not actually required. This will generally be
False on non-Windows platforms, unless optional functionality like
autoreset has been requested using kwargs to init()
'''
"""
return self.convert or self.strip or self.autoreset

def get_win32_calls(self):
if self.convert and winterm:
return {
AnsiStyle.RESET_ALL: (winterm.reset_all, ),
AnsiStyle.RESET_ALL: (winterm.reset_all,),
AnsiStyle.BRIGHT: (winterm.style, WinStyle.BRIGHT),
AnsiStyle.DIM: (winterm.style, WinStyle.NORMAL),
AnsiStyle.NORMAL: (winterm.style, WinStyle.NORMAL),
Expand All @@ -143,7 +154,7 @@ def get_win32_calls(self):
AnsiFore.MAGENTA: (winterm.fore, WinColor.MAGENTA),
AnsiFore.CYAN: (winterm.fore, WinColor.CYAN),
AnsiFore.WHITE: (winterm.fore, WinColor.GREY),
AnsiFore.RESET: (winterm.fore, ),
AnsiFore.RESET: (winterm.fore,),
AnsiFore.LIGHTBLACK_EX: (winterm.fore, WinColor.BLACK, True),
AnsiFore.LIGHTRED_EX: (winterm.fore, WinColor.RED, True),
AnsiFore.LIGHTGREEN_EX: (winterm.fore, WinColor.GREEN, True),
Expand All @@ -160,7 +171,7 @@ def get_win32_calls(self):
AnsiBack.MAGENTA: (winterm.back, WinColor.MAGENTA),
AnsiBack.CYAN: (winterm.back, WinColor.CYAN),
AnsiBack.WHITE: (winterm.back, WinColor.GREY),
AnsiBack.RESET: (winterm.back, ),
AnsiBack.RESET: (winterm.back,),
AnsiBack.LIGHTBLACK_EX: (winterm.back, WinColor.BLACK, True),
AnsiBack.LIGHTRED_EX: (winterm.back, WinColor.RED, True),
AnsiBack.LIGHTGREEN_EX: (winterm.back, WinColor.GREEN, True),
Expand All @@ -181,20 +192,18 @@ def write(self, text):
if self.autoreset:
self.reset_all()


def reset_all(self):
if self.convert:
self.call_win32('m', (0,))
self.call_win32("m", (0,))
elif not self.strip and not self.stream.closed:
self.wrapped.write(Style.RESET_ALL)


def write_and_convert(self, text):
'''
"""
Write the given text to our wrapped stream, stripping any ANSI
sequences from the text, and optionally converting them into win32
calls.
'''
"""
cursor = 0
text = self.convert_osc(text)
for match in self.ANSI_CSI_RE.finditer(text):
Expand All @@ -204,59 +213,54 @@ def write_and_convert(self, text):
cursor = end
self.write_plain_text(text, cursor, len(text))


def write_plain_text(self, text, start, end):
if start < end:
self.wrapped.write(text[start:end])
self.wrapped.flush()


def convert_ansi(self, paramstring, command):
if self.convert:
params = self.extract_params(command, paramstring)
self.call_win32(command, params)


def extract_params(self, command, paramstring):
if command in 'Hf':
params = tuple(int(p) if len(p) != 0 else 1 for p in paramstring.split(';'))
if command in "Hf":
params = tuple(int(p) if len(p) != 0 else 1 for p in paramstring.split(";"))
while len(params) < 2:
# defaults:
params = params + (1,)
else:
params = tuple(int(p) for p in paramstring.split(';') if len(p) != 0)
params = tuple(int(p) for p in paramstring.split(";") if len(p) != 0)
if len(params) == 0:
# defaults:
if command in 'JKm':
if command in "JKm":
params = (0,)
elif command in 'ABCD':
elif command in "ABCD":
params = (1,)

return params


def call_win32(self, command, params):
if command == 'm':
if command == "m":
for param in params:
if param in self.win32_calls:
func_args = self.win32_calls[param]
func = func_args[0]
args = func_args[1:]
kwargs = dict(on_stderr=self.on_stderr)
func(*args, **kwargs)
elif command in 'J':
elif command in "J":
winterm.erase_screen(params[0], on_stderr=self.on_stderr)
elif command in 'K':
elif command in "K":
winterm.erase_line(params[0], on_stderr=self.on_stderr)
elif command in 'Hf': # cursor position - absolute
elif command in "Hf": # cursor position - absolute
winterm.set_cursor_position(params, on_stderr=self.on_stderr)
elif command in 'ABCD': # cursor position - relative
elif command in "ABCD": # cursor position - relative
n = params[0]
# A - up, B - down, C - forward, D - back
x, y = {'A': (0, -n), 'B': (0, n), 'C': (n, 0), 'D': (-n, 0)}[command]
x, y = {"A": (0, -n), "B": (0, n), "C": (n, 0), "D": (-n, 0)}[command]
winterm.cursor_adjust(x, y, on_stderr=self.on_stderr)


def convert_osc(self, text):
for match in self.ANSI_OSC_RE.finditer(text):
start, end = match.span()
Expand All @@ -268,10 +272,9 @@ def convert_osc(self, text):
# 0 - change title and icon (we will only change title)
# 1 - change icon (we don't support this)
# 2 - change title
if params[0] in '02':
if params[0] in "02":
winterm.set_title(params[1])
return text


def flush(self):
self.wrapped.flush()

0 comments on commit d2a35c6

Please sign in to comment.