codekingpro/portable-devtools

1import re2import sys3from contextlib import suppress4from typing import Iterable, NamedTuple, Optional5 6from .color import Color7from .style import Style8from .text import Text9 10re_ansi = re.compile(11    r"""12(?:\x1b\](.*?)\x1b\\)|13(?:\x1b([(@-Z\\-_]|\[[0-?]*[ -/]*[@-~]))14""",15    re.VERBOSE,16)17 18 19class _AnsiToken(NamedTuple):20    """Result of ansi tokenized string."""21 22    plain: str = ""23    sgr: Optional[str] = ""24    osc: Optional[str] = ""25 26 27def _ansi_tokenize(ansi_text: str) -> Iterable[_AnsiToken]:28    """Tokenize a string in to plain text and ANSI codes.29 30    Args:31        ansi_text (str): A String containing ANSI codes.32 33    Yields:34        AnsiToken: A named tuple of (plain, sgr, osc)35    """36 37    position = 038    sgr: Optional[str]39    osc: Optional[str]40    for match in re_ansi.finditer(ansi_text):41        start, end = match.span(0)42        osc, sgr = match.groups()43        if start > position:44            yield _AnsiToken(ansi_text[position:start])45        if sgr:46            if sgr == "(":47                position = end + 148                continue49            if sgr.endswith("m"):50                yield _AnsiToken("", sgr[1:-1], osc)51        else:52            yield _AnsiToken("", sgr, osc)53        position = end54    if position < len(ansi_text):55        yield _AnsiToken(ansi_text[position:])56 57 58SGR_STYLE_MAP = {59    1: "bold",60    2: "dim",61    3: "italic",62    4: "underline",63    5: "blink",64    6: "blink2",65    7: "reverse",66    8: "conceal",67    9: "strike",68    21: "underline2",69    22: "not dim not bold",70    23: "not italic",71    24: "not underline",72    25: "not blink",73    26: "not blink2",74    27: "not reverse",75    28: "not conceal",76    29: "not strike",77    30: "color(0)",78    31: "color(1)",79    32: "color(2)",80    33: "color(3)",81    34: "color(4)",82    35: "color(5)",83    36: "color(6)",84    37: "color(7)",85    39: "default",86    40: "on color(0)",87    41: "on color(1)",88    42: "on color(2)",89    43: "on color(3)",90    44: "on color(4)",91    45: "on color(5)",92    46: "on color(6)",93    47: "on color(7)",94    49: "on default",95    51: "frame",96    52: "encircle",97    53: "overline",98    54: "not frame not encircle",99    55: "not overline",100    90: "color(8)",101    91: "color(9)",102    92: "color(10)",103    93: "color(11)",104    94: "color(12)",105    95: "color(13)",106    96: "color(14)",107    97: "color(15)",108    100: "on color(8)",109    101: "on color(9)",110    102: "on color(10)",111    103: "on color(11)",112    104: "on color(12)",113    105: "on color(13)",114    106: "on color(14)",115    107: "on color(15)",116}117 118 119class AnsiDecoder:120    """Translate ANSI code in to styled Text."""121 122    def __init__(self) -> None:123        self.style = Style.null()124 125    def decode(self, terminal_text: str) -> Iterable[Text]:126        """Decode ANSI codes in an iterable of lines.127 128        Args:129            lines (Iterable[str]): An iterable of lines of terminal output.130 131        Yields:132            Text: Marked up Text.133        """134        for line in terminal_text.splitlines():135            yield self.decode_line(line)136 137    def decode_line(self, line: str) -> Text:138        """Decode a line containing ansi codes.139 140        Args:141            line (str): A line of terminal output.142 143        Returns:144            Text: A Text instance marked up according to ansi codes.145        """146        from_ansi = Color.from_ansi147        from_rgb = Color.from_rgb148        _Style = Style149        text = Text()150        append = text.append151        line = line.rsplit("\r", 1)[-1]152        for plain_text, sgr, osc in _ansi_tokenize(line):153            if plain_text:154                append(plain_text, self.style or None)155            elif osc is not None:156                if osc.startswith("8;"):157                    _params, semicolon, link = osc[2:].partition(";")158                    if semicolon:159                        self.style = self.style.update_link(link or None)160            elif sgr is not None:161                # Translate in to semi-colon separated codes162                # Ignore invalid codes, because we want to be lenient163                codes = [164                    min(255, int(_code) if _code else 0)165                    for _code in sgr.split(";")166                    if _code.isdigit() or _code == ""167                ]168                iter_codes = iter(codes)169                for code in iter_codes:170                    if code == 0:171                        # reset172                        self.style = _Style.null()173                    elif code in SGR_STYLE_MAP:174                        # styles175                        self.style += _Style.parse(SGR_STYLE_MAP[code])176                    elif code == 38:177                        #  Foreground178                        with suppress(StopIteration):179                            color_type = next(iter_codes)180                            if color_type == 5:181                                self.style += _Style.from_color(182                                    from_ansi(next(iter_codes))183                                )184                            elif color_type == 2:185                                self.style += _Style.from_color(186                                    from_rgb(187                                        next(iter_codes),188                                        next(iter_codes),189                                        next(iter_codes),190                                    )191                                )192                    elif code == 48:193                        # Background194                        with suppress(StopIteration):195                            color_type = next(iter_codes)196                            if color_type == 5:197                                self.style += _Style.from_color(198                                    None, from_ansi(next(iter_codes))199                                )200                            elif color_type == 2:201                                self.style += _Style.from_color(202                                    None,203                                    from_rgb(204                                        next(iter_codes),205                                        next(iter_codes),206                                        next(iter_codes),207                                    ),208                                )209 210        return text211 212 213if sys.platform != "win32" and __name__ == "__main__":  # pragma: no cover214    import io215    import os216    import pty217    import sys218 219    decoder = AnsiDecoder()220 221    stdout = io.BytesIO()222 223    def read(fd: int) -> bytes:224        data = os.read(fd, 1024)225        stdout.write(data)226        return data227 228    pty.spawn(sys.argv[1:], read)229 230    from .console import Console231 232    console = Console(record=True)233 234    stdout_result = stdout.getvalue().decode("utf-8")235    print(stdout_result)236 237    for line in decoder.decode(stdout_result):238        console.print(line)239 240    console.save_html("stdout.html")241