1from __future__ import annotations 
    2 
    3from string import Formatter 
    4from typing import Generator 
    5 
    6from prompt_toolkit.output.vt100 import BG_ANSI_COLORS, FG_ANSI_COLORS 
    7from prompt_toolkit.output.vt100 import _256_colors as _256_colors_table 
    8 
    9from .base import StyleAndTextTuples 
    10 
    11__all__ = [ 
    12    "ANSI", 
    13    "ansi_escape", 
    14] 
    15 
    16 
    17class ANSI: 
    18    """ 
    19    ANSI formatted text. 
    20    Take something ANSI escaped text, for use as a formatted string. E.g. 
    21 
    22    :: 
    23 
    24        ANSI('\\x1b[31mhello \\x1b[32mworld') 
    25 
    26    Characters between ``\\001`` and ``\\002`` are supposed to have a zero width 
    27    when printed, but these are literally sent to the terminal output. This can 
    28    be used for instance, for inserting Final Term prompt commands.  They will 
    29    be translated into a prompt_toolkit '[ZeroWidthEscape]' fragment. 
    30    """ 
    31 
    32    def __init__(self, value: str) -> None: 
    33        self.value = value 
    34        self._formatted_text: StyleAndTextTuples = [] 
    35 
    36        # Default style attributes. 
    37        self._color: str | None = None 
    38        self._bgcolor: str | None = None 
    39        self._bold = False 
    40        self._dim = False 
    41        self._underline = False 
    42        self._strike = False 
    43        self._italic = False 
    44        self._blink = False 
    45        self._reverse = False 
    46        self._hidden = False 
    47 
    48        # Process received text. 
    49        parser = self._parse_corot() 
    50        parser.send(None)  # type: ignore 
    51        for c in value: 
    52            parser.send(c) 
    53 
    54    def _parse_corot(self) -> Generator[None, str, None]: 
    55        """ 
    56        Coroutine that parses the ANSI escape sequences. 
    57        """ 
    58        style = "" 
    59        formatted_text = self._formatted_text 
    60 
    61        while True: 
    62            # NOTE: CSI is a special token within a stream of characters that 
    63            #       introduces an ANSI control sequence used to set the 
    64            #       style attributes of the following characters. 
    65            csi = False 
    66 
    67            c = yield 
    68 
    69            # Everything between \001 and \002 should become a ZeroWidthEscape. 
    70            if c == "\001": 
    71                escaped_text = "" 
    72                while c != "\002": 
    73                    c = yield 
    74                    if c == "\002": 
    75                        formatted_text.append(("[ZeroWidthEscape]", escaped_text)) 
    76                        c = yield 
    77                        break 
    78                    else: 
    79                        escaped_text += c 
    80 
    81            # Check for CSI 
    82            if c == "\x1b": 
    83                # Start of color escape sequence. 
    84                square_bracket = yield 
    85                if square_bracket == "[": 
    86                    csi = True 
    87                else: 
    88                    continue 
    89            elif c == "\x9b": 
    90                csi = True 
    91 
    92            if csi: 
    93                # Got a CSI sequence. Color codes are following. 
    94                current = "" 
    95                params = [] 
    96 
    97                while True: 
    98                    char = yield 
    99 
    100                    # Construct number 
    101                    if char.isdigit(): 
    102                        current += char 
    103 
    104                    # Eval number 
    105                    else: 
    106                        # Limit and save number value 
    107                        params.append(min(int(current or 0), 9999)) 
    108 
    109                        # Get delimiter token if present 
    110                        if char == ";": 
    111                            current = "" 
    112 
    113                        # Check and evaluate color codes 
    114                        elif char == "m": 
    115                            # Set attributes and token. 
    116                            self._select_graphic_rendition(params) 
    117                            style = self._create_style_string() 
    118                            break 
    119 
    120                        # Check and evaluate cursor forward 
    121                        elif char == "C": 
    122                            for i in range(params[0]): 
    123                                # add <SPACE> using current style 
    124                                formatted_text.append((style, " ")) 
    125                            break 
    126 
    127                        else: 
    128                            # Ignore unsupported sequence. 
    129                            break 
    130            else: 
    131                # Add current character. 
    132                # NOTE: At this point, we could merge the current character 
    133                #       into the previous tuple if the style did not change, 
    134                #       however, it's not worth the effort given that it will 
    135                #       be "Exploded" once again when it's rendered to the 
    136                #       output. 
    137                formatted_text.append((style, c)) 
    138 
    139    def _select_graphic_rendition(self, attrs: list[int]) -> None: 
    140        """ 
    141        Taken a list of graphics attributes and apply changes. 
    142        """ 
    143        if not attrs: 
    144            attrs = [0] 
    145        else: 
    146            attrs = list(attrs[::-1]) 
    147 
    148        while attrs: 
    149            attr = attrs.pop() 
    150 
    151            if attr in _fg_colors: 
    152                self._color = _fg_colors[attr] 
    153            elif attr in _bg_colors: 
    154                self._bgcolor = _bg_colors[attr] 
    155            elif attr == 1: 
    156                self._bold = True 
    157            elif attr == 2: 
    158                self._dim = True 
    159            elif attr == 3: 
    160                self._italic = True 
    161            elif attr == 4: 
    162                self._underline = True 
    163            elif attr == 5: 
    164                self._blink = True  # Slow blink 
    165            elif attr == 6: 
    166                self._blink = True  # Fast blink 
    167            elif attr == 7: 
    168                self._reverse = True 
    169            elif attr == 8: 
    170                self._hidden = True 
    171            elif attr == 9: 
    172                self._strike = True 
    173            elif attr == 22: 
    174                self._bold = False  # Normal intensity 
    175                self._dim = False 
    176            elif attr == 23: 
    177                self._italic = False 
    178            elif attr == 24: 
    179                self._underline = False 
    180            elif attr == 25: 
    181                self._blink = False 
    182            elif attr == 27: 
    183                self._reverse = False 
    184            elif attr == 28: 
    185                self._hidden = False 
    186            elif attr == 29: 
    187                self._strike = False 
    188            elif not attr: 
    189                # Reset all style attributes 
    190                self._color = None 
    191                self._bgcolor = None 
    192                self._bold = False 
    193                self._dim = False 
    194                self._underline = False 
    195                self._strike = False 
    196                self._italic = False 
    197                self._blink = False 
    198                self._reverse = False 
    199                self._hidden = False 
    200 
    201            elif attr in (38, 48) and len(attrs) > 1: 
    202                n = attrs.pop() 
    203 
    204                # 256 colors. 
    205                if n == 5 and len(attrs) >= 1: 
    206                    if attr == 38: 
    207                        m = attrs.pop() 
    208                        self._color = _256_colors.get(m) 
    209                    elif attr == 48: 
    210                        m = attrs.pop() 
    211                        self._bgcolor = _256_colors.get(m) 
    212 
    213                # True colors. 
    214                if n == 2 and len(attrs) >= 3: 
    215                    try: 
    216                        color_str = ( 
    217                            f"#{attrs.pop():02x}{attrs.pop():02x}{attrs.pop():02x}" 
    218                        ) 
    219                    except IndexError: 
    220                        pass 
    221                    else: 
    222                        if attr == 38: 
    223                            self._color = color_str 
    224                        elif attr == 48: 
    225                            self._bgcolor = color_str 
    226 
    227    def _create_style_string(self) -> str: 
    228        """ 
    229        Turn current style flags into a string for usage in a formatted text. 
    230        """ 
    231        result = [] 
    232        if self._color: 
    233            result.append(self._color) 
    234        if self._bgcolor: 
    235            result.append("bg:" + self._bgcolor) 
    236        if self._bold: 
    237            result.append("bold") 
    238        if self._dim: 
    239            result.append("dim") 
    240        if self._underline: 
    241            result.append("underline") 
    242        if self._strike: 
    243            result.append("strike") 
    244        if self._italic: 
    245            result.append("italic") 
    246        if self._blink: 
    247            result.append("blink") 
    248        if self._reverse: 
    249            result.append("reverse") 
    250        if self._hidden: 
    251            result.append("hidden") 
    252 
    253        return " ".join(result) 
    254 
    255    def __repr__(self) -> str: 
    256        return f"ANSI({self.value!r})" 
    257 
    258    def __pt_formatted_text__(self) -> StyleAndTextTuples: 
    259        return self._formatted_text 
    260 
    261    def format(self, *args: str, **kwargs: str) -> ANSI: 
    262        """ 
    263        Like `str.format`, but make sure that the arguments are properly 
    264        escaped. (No ANSI escapes can be injected.) 
    265        """ 
    266        return ANSI(FORMATTER.vformat(self.value, args, kwargs)) 
    267 
    268    def __mod__(self, value: object) -> ANSI: 
    269        """ 
    270        ANSI('<b>%s</b>') % value 
    271        """ 
    272        if not isinstance(value, tuple): 
    273            value = (value,) 
    274 
    275        value = tuple(ansi_escape(i) for i in value) 
    276        return ANSI(self.value % value) 
    277 
    278 
    279# Mapping of the ANSI color codes to their names. 
    280_fg_colors = {v: k for k, v in FG_ANSI_COLORS.items()} 
    281_bg_colors = {v: k for k, v in BG_ANSI_COLORS.items()} 
    282 
    283# Mapping of the escape codes for 256colors to their 'ffffff' value. 
    284_256_colors = {} 
    285 
    286for i, (r, g, b) in enumerate(_256_colors_table.colors): 
    287    _256_colors[i] = f"#{r:02x}{g:02x}{b:02x}" 
    288 
    289 
    290def ansi_escape(text: object) -> str: 
    291    """ 
    292    Replace characters with a special meaning. 
    293    """ 
    294    return str(text).replace("\x1b", "?").replace("\b", "?") 
    295 
    296 
    297class ANSIFormatter(Formatter): 
    298    def format_field(self, value: object, format_spec: str) -> str: 
    299        return ansi_escape(format(value, format_spec)) 
    300 
    301 
    302FORMATTER = ANSIFormatter()