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