Skip to content

Commit

Permalink
Enable the use of rich for presenting output
Browse files Browse the repository at this point in the history
This makes it possible to present output with rich markup, within the
constraints of our logging infrastructure.

Further, diagnostic errors can now by presented using rich, using their
own special "[present-diagnostic]" marker string, since those need to be
handled differently from regular log messages and passed directly
through to rich's console object, after an indentation wrapper.
  • Loading branch information
pradyunsg committed Dec 3, 2021
1 parent 7d53c7c commit 9664158
Show file tree
Hide file tree
Showing 4 changed files with 69 additions and 95 deletions.
4 changes: 1 addition & 3 deletions src/pip/_internal/cli/base_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,6 @@
from optparse import Values
from typing import Any, Callable, List, Optional, Tuple

from pip._vendor import rich

from pip._internal.cli import cmdoptions
from pip._internal.cli.command_context import CommandContextMixIn
from pip._internal.cli.parser import ConfigOptionParser, UpdatingDefaultsHelpFormatter
Expand Down Expand Up @@ -168,7 +166,7 @@ def exc_logging_wrapper(*args: Any) -> int:
assert isinstance(status, int)
return status
except DiagnosticPipError as exc:
rich.print(exc, file=sys.stderr)
logger.error("[present-diagnostic]", exc)
logger.debug("Exception information:", exc_info=True)

return ERROR
Expand Down
6 changes: 3 additions & 3 deletions src/pip/_internal/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,9 @@ def _prefix_with_indent(
else:
text = console.render_str(s)

lines = text.wrap(console, console.width - width_offset)

return console.render_str(prefix) + console.render_str(f"\n{indent}").join(lines)
return console.render_str(prefix, overflow="ignore") + console.render_str(
f"\n{indent}", overflow="ignore"
).join(text.split(allow_blank=True))


class PipError(Exception):
Expand Down
146 changes: 61 additions & 85 deletions src/pip/_internal/utils/logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,28 +4,27 @@
import logging.handlers
import os
import sys
import threading
from dataclasses import dataclass
from logging import Filter
from typing import IO, Any, Callable, Iterator, Optional, TextIO, Type, cast

from typing import IO, Any, Iterator, Optional, TextIO, Type

from pip._vendor.rich.console import (
Console,
ConsoleOptions,
ConsoleRenderable,
RenderResult,
)
from pip._vendor.rich.highlighter import NullHighlighter
from pip._vendor.rich.logging import RichHandler
from pip._vendor.rich.segment import Segment

from pip._internal.exceptions import DiagnosticPipError
from pip._internal.utils._log import VERBOSE, getLogger
from pip._internal.utils.compat import WINDOWS
from pip._internal.utils.deprecation import DEPRECATION_MSG_PREFIX
from pip._internal.utils.misc import ensure_dir

try:
import threading
except ImportError:
import dummy_threading as threading # type: ignore


try:
from pip._vendor import colorama
# Lots of different errors can come from this, including SystemError and
# ImportError.
except Exception:
colorama = None


_log_state = threading.local()
subprocess_logger = getLogger("pip.subprocessor")

Expand Down Expand Up @@ -119,78 +118,54 @@ def format(self, record: logging.LogRecord) -> str:
return formatted


def _color_wrap(*colors: str) -> Callable[[str], str]:
def wrapped(inp: str) -> str:
return "".join(list(colors) + [inp, colorama.Style.RESET_ALL])

return wrapped


class ColorizedStreamHandler(logging.StreamHandler):

# Don't build up a list of colors if we don't have colorama
if colorama:
COLORS = [
# This needs to be in order from highest logging level to lowest.
(logging.ERROR, _color_wrap(colorama.Fore.RED)),
(logging.WARNING, _color_wrap(colorama.Fore.YELLOW)),
]
else:
COLORS = []

def __init__(self, stream: Optional[TextIO] = None, no_color: bool = None) -> None:
super().__init__(stream)
self._no_color = no_color

if WINDOWS and colorama:
self.stream = colorama.AnsiToWin32(self.stream)

def _using_stdout(self) -> bool:
"""
Return whether the handler is using sys.stdout.
"""
if WINDOWS and colorama:
# Then self.stream is an AnsiToWin32 object.
stream = cast(colorama.AnsiToWin32, self.stream)
return stream.wrapped is sys.stdout

return self.stream is sys.stdout

def should_color(self) -> bool:
# Don't colorize things if we do not have colorama or if told not to
if not colorama or self._no_color:
return False

real_stream = (
self.stream
if not isinstance(self.stream, colorama.AnsiToWin32)
else self.stream.wrapped
@dataclass
class IndentedRenderable:
renderable: ConsoleRenderable
indent: int

def __rich_console__(
self, console: Console, options: ConsoleOptions
) -> RenderResult:
segments = console.render(self.renderable, options)
lines = Segment.split_lines(segments)
for line in lines:
yield Segment(" " * self.indent)
yield from line
yield Segment("\n")


class RichPipStreamHandler(RichHandler):
def __init__(self, stream: Optional[TextIO], no_color: bool) -> None:
super().__init__(
console=Console(file=stream, no_color=no_color, soft_wrap=True),
show_time=False,
show_level=False,
show_path=False,
highlighter=NullHighlighter(),
)
self.KEYWORDS = []

# If the stream is a tty we should color it
if hasattr(real_stream, "isatty") and real_stream.isatty():
return True

# If we have an ANSI term we should color it
if os.environ.get("TERM") == "ANSI":
return True

# If anything else we should not color it
return False

def format(self, record: logging.LogRecord) -> str:
msg = super().format(record)

if self.should_color():
for level, color in self.COLORS:
if record.levelno >= level:
msg = color(msg)
break
# Our custom override on rich's logger, to make things work as we need them to.
def emit(self, record: logging.LogRecord) -> None:
# If we are given a diagnostic error to present, present it with indentation.
if (
record.msg == "[present-diagnostic]"
and len(record.args) == 1
and isinstance(record.args[0], DiagnosticPipError)
):
renderable = IndentedRenderable(record.args[0], indent=get_indentation())
else:
message = self.format(record)
renderable = self.render_message(record, message)

return msg
try:
self.console.print(renderable, overflow="ignore", crop=False)
except Exception:
self.handleError(record)

# The logging module says handleError() can be customized.
def handleError(self, record: logging.LogRecord) -> None:
"""Called when logging is unable to log some output."""

exc_class, exc = sys.exc_info()[:2]
# If a broken pipe occurred while calling write() or flush() on the
# stdout stream in logging's Handler.emit(), then raise our special
Expand All @@ -199,7 +174,7 @@ def handleError(self, record: logging.LogRecord) -> None:
if (
exc_class
and exc
and self._using_stdout()
and self.console.file is sys.stdout
and _is_broken_pipe_error(exc_class, exc)
):
raise BrokenStdoutLoggingError()
Expand Down Expand Up @@ -275,7 +250,8 @@ def setup_logging(verbosity: int, no_color: bool, user_log_file: Optional[str])
"stderr": "ext://sys.stderr",
}
handler_classes = {
"stream": "pip._internal.utils.logging.ColorizedStreamHandler",
"subprocess": "logging.StreamHandler",
"stream": "pip._internal.utils.logging.RichPipStreamHandler",
"file": "pip._internal.utils.logging.BetterRotatingFileHandler",
}
handlers = ["console", "console_errors", "console_subprocess"] + (
Expand Down
8 changes: 4 additions & 4 deletions tests/unit/test_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@

from pip._internal.utils.logging import (
BrokenStdoutLoggingError,
ColorizedStreamHandler,
IndentingFormatter,
RichPipStreamHandler,
indent_log,
)
from pip._internal.utils.misc import captured_stderr, captured_stdout
Expand Down Expand Up @@ -142,7 +142,7 @@ def test_broken_pipe_in_stderr_flush(self) -> None:
record = self._make_log_record()

with captured_stderr() as stderr:
handler = ColorizedStreamHandler(stream=stderr)
handler = RichPipStreamHandler(stream=stderr, no_color=True)
with patch("sys.stderr.flush") as mock_flush:
mock_flush.side_effect = BrokenPipeError()
# The emit() call raises no exception.
Expand All @@ -165,7 +165,7 @@ def test_broken_pipe_in_stdout_write(self) -> None:
record = self._make_log_record()

with captured_stdout() as stdout:
handler = ColorizedStreamHandler(stream=stdout)
handler = RichPipStreamHandler(stream=stdout, no_color=True)
with patch("sys.stdout.write") as mock_write:
mock_write.side_effect = BrokenPipeError()
with pytest.raises(BrokenStdoutLoggingError):
Expand All @@ -180,7 +180,7 @@ def test_broken_pipe_in_stdout_flush(self) -> None:
record = self._make_log_record()

with captured_stdout() as stdout:
handler = ColorizedStreamHandler(stream=stdout)
handler = RichPipStreamHandler(stream=stdout, no_color=True)
with patch("sys.stdout.flush") as mock_flush:
mock_flush.side_effect = BrokenPipeError()
with pytest.raises(BrokenStdoutLoggingError):
Expand Down

0 comments on commit 9664158

Please sign in to comment.