Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/prompt_toolkit/formatted_text/ansi.py: 77%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

184 statements  

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