Merge pull request #431 from willmcgugan/ansi-decoder

Ansi decoder
This commit is contained in:
Will McGugan 2020-11-07 15:30:41 +00:00 committed by GitHub
commit 4508ed3a43
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
13 changed files with 368 additions and 30 deletions

View File

@ -12,6 +12,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- Added tracebacks_show_locals parameter to RichHandler
- Applied dim=True to indent guide styles
- Added max_string to Pretty
- Added rich.ansi.AnsiDecoder
- Added decoding of ansi codes to captured stdout in Progress
- Hid progress bars from html export
## [9.1.0] - 2020-10-23

View File

@ -18,6 +18,7 @@ Reference
reference/measure.rst
reference/padding.rst
reference/panel.rst
reference/pretty.rst
reference/progress.rst
reference/progress_bar.rst
reference/prompt.rst

View File

@ -0,0 +1,6 @@
rich.pretty
===========
.. automodule:: rich.pretty
:members:

229
rich/ansi.py Normal file
View File

@ -0,0 +1,229 @@
from contextlib import suppress
import re
from typing import Iterable, NamedTuple
from .color import Color
from .style import Style
from .text import Text
re_ansi = re.compile(r"(?:\x1b\[(.*?)m)|(?:\x1b\](.*?)\x1b\\)")
re_csi = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
class _AnsiToken(NamedTuple):
"""Result of ansi tokenized string."""
plain: str = ""
sgr: str = ""
osc: str = ""
def _ansi_tokenize(ansi_text: str) -> Iterable[_AnsiToken]:
"""Tokenize a string in to plain text and ANSI codes.
Args:
ansi_text (str): A String containing ANSI codes.
Yields:
AnsiToken: A named tuple of (plain, sgr, osc)
"""
def remove_csi(ansi_text: str) -> str:
"""Remove unknown CSI sequences."""
return re_csi.sub("", ansi_text)
position = 0
for match in re_ansi.finditer(ansi_text):
start, end = match.span(0)
sgr, osc = match.groups()
if start > position:
yield _AnsiToken(remove_csi(ansi_text[position:start]))
yield _AnsiToken("", sgr, osc)
position = end
if position < len(ansi_text):
yield _AnsiToken(remove_csi(ansi_text[position:]))
SGR_STYLE_MAP = {
1: "bold",
2: "dim",
3: "italic",
4: "underline",
5: "blink",
6: "blink2",
7: "reverse",
8: "conceal",
9: "strike",
21: "underline2",
22: "not dim not bold",
23: "not italic",
24: "not underline",
25: "not blink",
26: "not blink2",
27: "not reverse",
28: "not conceal",
29: "not strike",
30: "color(0)",
31: "color(1)",
32: "color(2)",
33: "color(3)",
34: "color(4)",
35: "color(5)",
36: "color(6)",
37: "color(7)",
39: "default",
40: "on color(0)",
41: "on color(1)",
42: "on color(2)",
43: "on color(3)",
44: "on color(4)",
45: "on color(5)",
46: "on color(6)",
47: "on color(7)",
49: "on default",
51: "frame",
52: "encircle",
53: "overline",
54: "not frame not encircle",
55: "not overline",
90: "color(8)",
91: "color(9)",
92: "color(10)",
93: "color(11)",
94: "color(12)",
95: "color(13)",
96: "color(14)",
97: "color(15)",
100: "on color(8)",
101: "on color(9)",
102: "on color(10)",
103: "on color(11)",
104: "on color(12)",
105: "on color(13)",
106: "on color(14)",
107: "on color(15)",
}
class AnsiDecoder:
"""Translate ANSI code in to styled Text."""
def __init__(self) -> None:
self.style = Style.null()
def decode(self, terminal_text: str) -> Iterable[Text]:
"""Decode ANSI codes in an interable of lines.
Args:
lines (Iterable[str]): An iterable of lines of terminal output.
Yields:
Text: Marked up Text.
"""
for line in terminal_text.splitlines():
yield self.decode_line(line)
def decode_line(self, line: str) -> Text:
"""Decode a line containing ansi codes.
Args:
line (str): A line of terminal output.
Returns:
Text: A Text instance marked up according to ansi codes.
"""
from_ansi = Color.from_ansi
from_rgb = Color.from_rgb
_Style = Style
text = Text()
append = text.append
line = line.rsplit("\r", 1)[-1]
for token in _ansi_tokenize(line):
plain_text, sgr, osc = token
if plain_text:
append(plain_text, self.style or None)
elif osc:
if osc.startswith("8;"):
_params, semicolon, link = osc[2:].partition(";")
if semicolon:
self.style = self.style.update_link(link or None)
elif sgr:
# Translate in to semi-colon separated codes
# Ignore invalid codes, because we want to be lenient
codes = [
min(255, int(_code)) for _code in sgr.split(";") if _code.isdigit()
]
iter_codes = iter(codes)
for code in iter_codes:
if code == 0:
# reset
self.style = _Style.null()
elif code in SGR_STYLE_MAP:
# styles
self.style += _Style.parse(SGR_STYLE_MAP[code])
elif code == 38:
#  Foreground
with suppress(StopIteration):
color_type = next(iter_codes)
if color_type == 5:
self.style += _Style.from_color(
from_ansi(next(iter_codes))
)
elif color_type == 2:
self.style += _Style.from_color(
from_rgb(
next(iter_codes),
next(iter_codes),
next(iter_codes),
)
)
elif code == 48:
# Background
with suppress(StopIteration):
color_type = next(iter_codes)
if color_type == 5:
self.style += _Style.from_color(
None, from_ansi(next(iter_codes))
)
elif color_type == 2:
self.style += _Style.from_color(
None,
from_rgb(
next(iter_codes),
next(iter_codes),
next(iter_codes),
),
)
return text
if __name__ == "__main__": # pragma: no cover
import pty
import io
import os
import sys
decoder = AnsiDecoder()
stdout = io.BytesIO()
def read(fd):
data = os.read(fd, 1024)
stdout.write(data)
return data
pty.spawn(sys.argv[1:], read)
from .console import Console
console = Console(record=True)
stdout_result = stdout.getvalue().decode("utf-8")
print(stdout_result)
for line in decoder.decode(stdout_result):
console.print(line)
console.save_html("stdout.html")

View File

@ -327,6 +327,22 @@ class Color(NamedTuple):
assert self.number is None
return theme.foreground_color if foreground else theme.background_color
@classmethod
def from_ansi(cls, number: int) -> "Color":
"""Create a Color number from it's 8-bit ansi number.
Args:
number (int): A number between 0-255 inclusive.
Returns:
Color: A new Color instance.
"""
return cls(
name=f"color({number})",
type=(ColorType.STANDARD if number < 16 else ColorType.EIGHT_BIT),
number=number,
)
@classmethod
def from_triplet(cls, triplet: "ColorTriplet") -> "Color":
"""Create a truecolor RGB color from a triplet of values.

View File

@ -1274,7 +1274,7 @@ class Console:
self._record_buffer.extend(buffer)
not_terminal = not self.is_terminal
for text, style, is_control in buffer:
if style and not is_control:
if style:
append(
style.render(
text,

View File

@ -55,8 +55,8 @@ class LiveRender:
) -> RenderResult:
style = console.get_style(self.style)
lines = console.render_lines(self.renderable, options, style=style, pad=False)
shape = Segment.get_shape(lines)
_Segment = Segment
shape = _Segment.get_shape(lines)
if self._shape is None:
self._shape = shape
else:
@ -68,8 +68,8 @@ class LiveRender:
)
width, height = self._shape
lines = Segment.set_shape(lines, width, height)
lines = _Segment.set_shape(lines, width, height)
for last, line in loop_last(lines):
yield from line
yield from _Segment.make_control(line)
if not last:
yield Segment.line()
yield _Segment.line(is_control=True)

View File

@ -27,7 +27,7 @@ from typing import (
)
from . import filesize, get_console
from .progress_bar import ProgressBar
from .ansi import AnsiDecoder
from .console import (
Console,
ConsoleRenderable,
@ -40,6 +40,7 @@ from .control import Control
from .highlighter import Highlighter
from .jupyter import JupyterMixin
from .live_render import LiveRender
from .progress_bar import ProgressBar
from .style import StyleType
from .table import Table
from .text import Text
@ -480,6 +481,7 @@ class _FileProxy(io.TextIOBase):
self.__console = console
self.__file = file
self.__buffer: List[str] = []
self.__ansi_decoder = AnsiDecoder()
def __getattr__(self, name: str) -> Any:
return getattr(self.__file, name)
@ -498,7 +500,9 @@ class _FileProxy(io.TextIOBase):
if lines:
console = self.__console
with console:
output = "\n".join(lines)
output = Text("\n").join(
self.__ansi_decoder.decode_line(line) for line in lines
)
console.print(output, markup=False, emoji=False, highlight=False)
return len(text)
@ -846,8 +850,8 @@ class Progress(JupyterMixin, RenderHook):
"""Refresh (render) the progress information."""
if self.console.is_jupyter: # pragma: no cover
try:
from ipywidgets import Output
from IPython.display import display
from ipywidgets import Output
except ImportError:
import warnings
@ -974,13 +978,13 @@ class Progress(JupyterMixin, RenderHook):
if __name__ == "__main__": # pragma: no coverage
import time
import random
import time
from .panel import Panel
from .rule import Rule
from .syntax import Syntax
from .table import Table
from .rule import Rule
syntax = Syntax(
'''def loop_last(values: Iterable[T]) -> Iterable[Tuple[bool, T]]:
@ -1017,16 +1021,20 @@ if __name__ == "__main__": # pragma: no coverage
examples = cycle(progress_renderables)
console = Console()
with Progress(console=console, transient=True) as progress:
console = Console(record=True)
try:
with Progress(console=console, transient=True) as progress:
task1 = progress.add_task("[red]Downloading", total=1000)
task2 = progress.add_task("[green]Processing", total=1000)
task3 = progress.add_task("[yellow]Thinking", total=1000, start=False)
task1 = progress.add_task("[red]Downloading", total=1000)
task2 = progress.add_task("[green]Processing", total=1000)
task3 = progress.add_task("[yellow]Thinking", total=1000, start=False)
while not progress.finished:
progress.update(task1, advance=0.5)
progress.update(task2, advance=0.3)
time.sleep(0.01)
if random.randint(0, 100) < 1:
progress.log(next(examples))
while not progress.finished:
progress.update(task1, advance=0.5)
progress.update(task2, advance=0.3)
time.sleep(0.01)
if random.randint(0, 100) < 1:
progress.log(next(examples))
except:
console.save_html("progress.html")
print("wrote progress.html")

View File

@ -41,21 +41,31 @@ class Segment(NamedTuple):
return 0 if self.is_control else cell_len(self.text)
@classmethod
def control(cls, text: str) -> "Segment":
def control(cls, text: str, style: Optional[Style] = None) -> "Segment":
"""Create a Segment with control codes.
Args:
text (str): Text containing non-printable control codes.
style (Optional[style]): Optional style.
Returns:
Segment: A Segment instance with ``is_control=True``.
"""
return Segment(text, is_control=True)
return Segment(text, style, is_control=True)
@classmethod
def line(cls) -> "Segment":
def make_control(cls, segments: Iterable["Segment"]) -> Iterable["Segment"]:
"""Convert all segments in to control segments.
Returns:
Iterable[Segments]: Segments with is_control=True
"""
return [cls(text, style, True) for text, style, _ in segments]
@classmethod
def line(cls, is_control: bool = False) -> "Segment":
"""Make a new line segment."""
return Segment("\n")
return Segment("\n", is_control=is_control)
@classmethod
def apply_style(

View File

@ -175,6 +175,35 @@ class Style:
"""Create an 'null' style, equivalent to Style(), but more performant."""
return NULL_STYLE
@classmethod
def from_color(cls, color: Color = None, bgcolor: Color = None) -> "Style":
"""Create a new style with colors and no attributes.
Returns:
color (Optional[Color]): A (foreground) color, or None for no color. Defaults to None.
bgcolor (Optional[Color]): A (background) color, or None for no color. Defaults to None.
"""
style = cls.__new__(Style)
style._ansi = None
style._style_definition = None
style._color = color
style._bgcolor = bgcolor
style._set_attributes = 0
style._attributes = 0
style._link = None
style._link_id = ""
style._hash = hash(
(
color,
bgcolor,
None,
None,
None,
)
)
style._null = not (color or bgcolor)
return style
bold = _Bit(0)
dim = _Bit(1)
italic = _Bit(2)
@ -355,7 +384,7 @@ class Style:
return Style(bgcolor=self.bgcolor)
@classmethod
@lru_cache(maxsize=1024)
@lru_cache(maxsize=4096)
def parse(cls, style_definition: str) -> "Style":
"""Parse a style definition.
@ -369,7 +398,7 @@ class Style:
`Style`: A Style instance.
"""
if style_definition.strip() == "none":
return cls()
return cls.null()
style_attributes = {
"dim": "dim",

32
tests/test_ansi.py Normal file
View File

@ -0,0 +1,32 @@
import io
from rich.ansi import AnsiDecoder
from rich.console import Console
from rich.style import Style
from rich.text import Span, Text
def test_decode():
console = Console(
force_terminal=True, legacy_windows=False, color_system="truecolor"
)
console.begin_capture()
console.print("Hello")
console.print("[b]foo[/b]")
console.print("[link http://example.org]bar")
console.print("[#ff0000 on color(200)]red")
console.print("[color(200) on #ff0000]red")
terminal_codes = console.end_capture()
decoder = AnsiDecoder()
lines = list(decoder.decode(terminal_codes))
expected = [
Text("Hello"),
Text("foo", spans=[Span(0, 3, Style.parse("bold"))]),
Text("bar", spans=[Span(0, 3, Style.parse("link http://example.org"))]),
Text("red", spans=[Span(0, 3, Style.parse("#ff0000 on color(200)"))]),
Text("red", spans=[Span(0, 3, Style.parse("color(200) on #ff0000"))]),
]
assert lines == expected

View File

@ -73,6 +73,10 @@ def test_from_rgb() -> None:
)
def test_from_ansi() -> None:
assert Color.from_ansi(1) == Color("color(1)", ColorType.STANDARD, 1)
def test_default() -> None:
assert Color.default() == Color("default", ColorType.DEFAULT, None, None)

View File

@ -37,7 +37,7 @@ def test_rich_console(live_render):
encoding="utf-8",
)
rich_console = live_render.__rich_console__(Console(), options)
assert [Segment("my string", Style.parse("none"))] == list(rich_console)
assert [Segment.control("my string", Style.parse("none"))] == list(rich_console)
live_render.style = "red"
rich_console = live_render.__rich_console__(Console(), options)
assert [Segment("my string", Style.parse("red"))] == list(rich_console)
assert [Segment.control("my string", Style.parse("red"))] == list(rich_console)