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()