Skip to content

Fix --color being silently stripped on redirected/piped output #192

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
agu2347 wants to merge 1 commit into from
Closed

Fix --color being silently stripped on redirected/piped output #192

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
138 changes: 55 additions & 83 deletions graphtage/printer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,23 +9,25 @@
the command line).

Attributes:
DEFAULT_PRINTER (Printer): A default :class:`Printer` instance printing to :attr:`sys.stdout`. Read it through
:func:`get_default_printer` rather than importing the name, because :func:`set_default_printer` replaces it.

NULL_PRINTER (Printer): A :class:`Printer` instance that discards everything written to it.
DEFAULT_PRINTER (Printer): A default :class:`Printer` instance printing to :attr:`sys.stdout`.

"""

import logging
import os
import sys
from abc import abstractmethod
from collections import defaultdict
from functools import wraps
from typing import Any, Optional, Protocol, Union
from typing import Any, Dict, List, Optional, Set, Type, Union
if sys.version_info[0] < 3 or sys.version_info[1] < 7:
Protocol = object
else:
from typing_extensions import Protocol

import colorama
from colorama import Back, Fore, Style
from colorama.ansi import AnsiBack, AnsiFore, AnsiStyle
from colorama.ansi import AnsiFore, AnsiBack, AnsiStyle

from .progress import StatusWriter
from .version import VERSION_STRING
Expand Down Expand Up @@ -69,8 +71,8 @@ def raw_write(self, s: str) -> int:
raise NotImplementedError()


STRIKETHROUGH = '\u0336'
UNDER_PLUS = '\u031F'
STRIKETHROUGH = chr(0x0336) # combining long stroke overlay
UNDER_PLUS = chr(0x031F) # combining plus sign below


class CombiningMarkWriter(RawWriter):
Expand All @@ -85,7 +87,7 @@ def __init__(self, parent: RawWriter):
"""
self.parent: RawWriter = parent
"""This writer's parent."""
self._marks: set[str] = set()
self._marks: Set[str] = set()
self.enabled: bool = True
"""Whether or not combining marks will be added."""

Expand All @@ -102,7 +104,7 @@ def context(self, *combining_marks: str) -> 'CombiningMarkContext':
return CombiningMarkContext(self, *combining_marks)

@property
def marks(self) -> set[str]:
def marks(self) -> Set[str]:
"""Returns the set of combining marks in this writer."""
return self._marks

Expand Down Expand Up @@ -144,8 +146,8 @@ class CombiningMarkContext:
"""A context returned by :meth:`CombiningMarkWriter.context`."""
def __init__(self, writer: CombiningMarkWriter, *combining_marks: str):
self.writer: CombiningMarkWriter = writer
self.marks: set[str] = set(combining_marks)
self._state_before: set[str] | None = None
self.marks: Set[str] = set(combining_marks)
self._state_before: Optional[Set[str]] = None

def __enter__(self) -> CombiningMarkWriter:
self._state_before = set(self.writer.marks)
Expand All @@ -164,9 +166,9 @@ class ANSIContext:
def __init__(
self,
stream: Union[RawWriter, 'ANSIContext'],
fore: AnsiFore | None = None,
back: AnsiBack | None = None,
style: AnsiStyle | None = None,
fore: Optional[AnsiFore] = None,
back: Optional[AnsiBack] = None,
style: Optional[AnsiStyle] = None,
):
"""Initializes a context.

Expand All @@ -181,15 +183,15 @@ def __init__(
"""
if isinstance(stream, ANSIContext):
self.stream: RawWriter = stream.stream
self._parent: ANSIContext | None = stream
self._parent: Optional['ANSIContext'] = stream
else:
self.stream: RawWriter = stream
self._parent: ANSIContext | None = None
self._fore: AnsiFore | None = fore
self._back: AnsiBack | None = back
self._style: AnsiStyle | None = style
self._start_code: str | None = None
self._end_code: str | None = None
self._parent: Optional['ANSIContext'] = None
self._fore: Optional[AnsiFore] = fore
self._back: Optional[AnsiBack] = back
self._style: Optional[AnsiStyle] = style
self._start_code: Optional[str] = None
self._end_code: Optional[str] = None
self.is_applied: bool = False
"""Keeps track of whether this context's options have already been applied to the underlying stream."""

Expand Down Expand Up @@ -217,7 +219,7 @@ def _set_codes(self):
contexts = ANSI_CONTEXT_STACK[self.stream]
if contexts:
if self._parent is None:
self._parent: ANSIContext | None = contexts[-1]
self._parent: Optional['ANSIContext'] = contexts[-1]
else:
if not self.root.is_applied:
self.root._parent = contexts[-1]
Expand Down Expand Up @@ -250,23 +252,23 @@ def _set_codes(self):
self._end_code += parent_end_code

@property
def fore(self) -> AnsiFore | None:
def fore(self) -> Optional[AnsiFore]:
"""The computed foreground color of this context."""
if self._fore is None and self._parent is not None:
return self._parent.fore
else:
return self._fore

@property
def back(self) -> AnsiBack | None:
def back(self) -> Optional[AnsiBack]:
"""The computed background color of this context."""
if self._back is None and self._parent is not None:
return self._parent.back
else:
return self._back

@property
def style(self) -> AnsiStyle | None:
def style(self) -> Optional[AnsiStyle]:
"""The computed style of this context."""
if self._style is None and self._parent is not None:
return self._parent.style
Expand Down Expand Up @@ -361,7 +363,7 @@ def _set_codes(self):
contexts = ANSI_CONTEXT_STACK[self.stream]
if contexts:
if self._parent is None:
self._parent: ANSIContext | None = contexts[-1]
self._parent: Optional['ANSIContext'] = contexts[-1]
else:
if not self.root.is_applied:
self.root._parent = contexts[-1]
Expand All @@ -378,11 +380,11 @@ def _set_codes(self):
style += f"background-color: {self.get_back(self._back)};"
if self._style is not None and (self._parent is None or self._style != self.parent.style):
if self._style == Style.BRIGHT:
style += "font-weight: bold; opacity: 1.0;"
style += f"font-weight: bold; opacity: 1.0;"
elif self._style == Style.DIM:
style += "opacity: 0.6; font-weight: normal;"
style += f"opacity: 0.6; font-weight: normal;"
else:
style += "font-weight: normal; opacity: 1.0;"
style += f"font-weight: normal; opacity: 1.0;"

if style:
self._start_code = f'{self._start_code}<span style="{style}">'
Expand All @@ -393,7 +395,7 @@ def _set_codes(self):
self._end_code = f"{self._end_code}{parent_end_code}"


ONLY_ANSI_FUNCS: set[str] = set()
ONLY_ANSI_FUNCS: Set[str] = set()


def only_ansi(func):
Expand Down Expand Up @@ -437,36 +439,18 @@ def fake_fun(*args, **kwargs):
return getattr(self._printer, item)


ANSI_CONTEXT_STACK: dict[Writer, list[ANSIContext]] = defaultdict(list)


def enable_ansi_support(force_color: bool = False):
"""Prepares :attr:`sys.stdout` and :attr:`sys.stderr` to receive ANSI escape sequences.

On a legacy Windows console, :mod:`colorama` replaces both streams with wrappers that translate the escape
sequences into Win32 console calls. A :class:`Printer` captures its output stream when it is constructed, so call
this function first; a printer constructed beforehand writes past the wrapper and its color is lost.

This function mutates global state, so call it from an application entry point rather than from library code.

Args:
force_color: If :const:`True`, keep the escape sequences even when the output stream is not a terminal.
:mod:`colorama` strips them in that case by default, which would discard color that the user explicitly
requested.

"""
colorama.init(strip=False if force_color else None)
ANSI_CONTEXT_STACK: Dict[Writer, List[ANSIContext]] = defaultdict(list)


class Printer(StatusWriter, RawWriter):
"""An ANSI color and status printer."""

def __init__(
self,
out_stream: Writer | None = None,
ansi_color: bool | None = None,
out_stream: Optional[Writer] = None,
ansi_color: Optional[bool] = None,
quiet: bool = False,
options: dict[str, Any] | None = None
options: Optional[Dict[str, Any]] = None
):
"""Initializes a Printer.

Expand All @@ -485,7 +469,7 @@ def __init__(
out_stream=out_stream,
quiet=quiet
)
self._context_type: type[ANSIContext] = ANSIContext
self._context_type: Type[ANSIContext] = ANSIContext
self.out_stream: CombiningMarkWriter = CombiningMarkWriter(self)
"""The stream wrapped by this printer."""
self.indents: int = 0
Expand All @@ -494,6 +478,8 @@ def __init__(
"""The string used for each indent step (default is four spaces)."""
self._ansi_color = None
self.ansi_color = ansi_color
if self.ansi_color:
colorama.init()
self._strikethrough = False
self._plusthrough = False
if options is not None:
Expand All @@ -509,7 +495,7 @@ def ansi_color(self) -> bool:
return self._ansi_color

@ansi_color.setter
def ansi_color(self, is_color: bool | None):
def ansi_color(self, is_color: Optional[bool]):
if is_color is None:
self._ansi_color = self.out_stream.isatty()
else:
Expand Down Expand Up @@ -585,7 +571,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
class HTMLPrinter(Printer):
"""A Printer that outputs in HTML."""

def __init__(self, *args, title: str | None = None, **kwargs):
def __init__(self, *args, title: Optional[str] = None, **kwargs):
super().__init__(*args, **kwargs)
self._context_type = HTMLANSIContext
self.raw_write("<html>")
Expand Down Expand Up @@ -668,35 +654,21 @@ def write(self, s: str) -> int:
DEFAULT_PRINTER: Printer = Printer()


def get_default_printer() -> Printer:
"""Returns the printer that library code uses when the caller does not supply one.

Call this instead of importing :attr:`DEFAULT_PRINTER` by name. :func:`set_default_printer` rebinds the module
attribute, which a name bound by ``from .printer import DEFAULT_PRINTER`` never observes: such a name keeps
referring to the printer that was current when the importing module was first loaded.

Returns:
Printer: The printer most recently passed to :func:`set_default_printer`, or :attr:`DEFAULT_PRINTER` if that
function was never called.

"""
return DEFAULT_PRINTER


def set_default_printer(printer: Printer):
"""Installs :obj:`printer` as the printer returned by :func:`get_default_printer`.

This mutates global state, so call it from an application entry point rather than from library code.

Args:
printer: The printer to install.

class NullWriter(Writer):
"""A writer that discards everything written to it.

``isatty()`` must return :const:`False` here. :class:`Printer` defaults
``ansi_color`` to ``out_stream.isatty()`` and calls ``colorama.init()``
as a side effect whenever color ends up enabled (see
:meth:`Printer.__init__`). ``NULL_PRINTER`` below is constructed at
import time, so if this returned :const:`True`, importing this module
would unconditionally call ``colorama.init()``, which globally replaces
:attr:`sys.stdout`/:attr:`sys.stderr` with colorama's stripping wrapper
-- silently disabling ``--color`` for any later, real ``Printer`` that
writes to a redirected/piped stream (#128). A sink that discards every
write has no terminal to color in the first place.
"""
global DEFAULT_PRINTER
DEFAULT_PRINTER = printer


class NullWriter(Writer):
def write(self, s: str) -> int:
return 0

Expand Down
Loading