238 lines
6.7 KiB
Python
238 lines
6.7 KiB
Python
import re
|
||
import sys
|
||
from contextlib import suppress
|
||
from typing import Iterable, NamedTuple, Optional
|
||
|
||
from .color import Color
|
||
from .style import Style
|
||
from .text import Text
|
||
|
||
re_ansi = re.compile(
|
||
r"""
|
||
(?:\x1b\](.*?)\x1b\\)|
|
||
(?:\x1b([(@-Z\\-_]|\[[0-?]*[ -/]*[@-~]))
|
||
""",
|
||
re.VERBOSE,
|
||
)
|
||
|
||
|
||
class _AnsiToken(NamedTuple):
|
||
"""Result of ansi tokenized string."""
|
||
|
||
plain: str = ""
|
||
sgr: Optional[str] = ""
|
||
osc: Optional[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)
|
||
"""
|
||
|
||
position = 0
|
||
sgr: Optional[str]
|
||
osc: Optional[str]
|
||
for match in re_ansi.finditer(ansi_text):
|
||
start, end = match.span(0)
|
||
osc, sgr = match.groups()
|
||
if start > position:
|
||
yield _AnsiToken(ansi_text[position:start])
|
||
if sgr:
|
||
if sgr.endswith("m"):
|
||
yield _AnsiToken("", sgr[1:-1], osc)
|
||
else:
|
||
yield _AnsiToken("", sgr, osc)
|
||
position = end
|
||
if position < len(ansi_text):
|
||
yield _AnsiToken(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 plain_text, sgr, osc in _ansi_tokenize(line):
|
||
if plain_text:
|
||
append(plain_text, self.style or None)
|
||
elif osc is not None:
|
||
if osc.startswith("8;"):
|
||
_params, semicolon, link = osc[2:].partition(";")
|
||
if semicolon:
|
||
self.style = self.style.update_link(link or None)
|
||
elif sgr is not None:
|
||
# Translate in to semi-colon separated codes
|
||
# Ignore invalid codes, because we want to be lenient
|
||
codes = [
|
||
min(255, int(_code) if _code else 0)
|
||
for _code in sgr.split(";")
|
||
if _code.isdigit() or _code == ""
|
||
]
|
||
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 sys.platform != "win32" and __name__ == "__main__": # pragma: no cover
|
||
import io
|
||
import os
|
||
import pty
|
||
import sys
|
||
|
||
decoder = AnsiDecoder()
|
||
|
||
stdout = io.BytesIO()
|
||
|
||
def read(fd: int) -> bytes:
|
||
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")
|