sgr_state.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338
  1. """
  2. SGR (Select Graphic Rendition) state tracking for terminal escape sequences.
  3. This module provides functions for tracking and propagating terminal styling (bold, italic, colors,
  4. etc.) via public API propagate_sgr(), and its dependent functions, cut() and wrap(). It only has
  5. attributes necessary to perform its functions, eg 'RED' and 'BLUE' attributes are not defined.
  6. """
  7. from __future__ import annotations
  8. # std imports
  9. import re
  10. from enum import IntEnum
  11. from typing import TYPE_CHECKING, Iterator, NamedTuple
  12. if TYPE_CHECKING: # pragma: no cover
  13. from typing import Sequence
  14. class _SGR(IntEnum):
  15. """
  16. SGR (Select Graphic Rendition) parameter codes.
  17. References:
  18. - https://invisible-island.net/xterm/ctlseqs/ctlseqs.html
  19. - https://github.com/tehmaze/ansi/tree/master/ansi/colour
  20. """
  21. RESET = 0
  22. BOLD = 1
  23. DIM = 2
  24. ITALIC = 3
  25. UNDERLINE = 4
  26. BLINK = 5
  27. RAPID_BLINK = 6
  28. INVERSE = 7
  29. HIDDEN = 8
  30. STRIKETHROUGH = 9
  31. DOUBLE_UNDERLINE = 21
  32. BOLD_DIM_OFF = 22
  33. ITALIC_OFF = 23
  34. UNDERLINE_OFF = 24
  35. BLINK_OFF = 25
  36. INVERSE_OFF = 27
  37. HIDDEN_OFF = 28
  38. STRIKETHROUGH_OFF = 29
  39. FG_BLACK = 30
  40. FG_WHITE = 37
  41. FG_EXTENDED = 38
  42. FG_DEFAULT = 39
  43. BG_BLACK = 40
  44. BG_WHITE = 47
  45. BG_EXTENDED = 48
  46. BG_DEFAULT = 49
  47. FG_BRIGHT_BLACK = 90
  48. FG_BRIGHT_WHITE = 97
  49. BG_BRIGHT_BLACK = 100
  50. BG_BRIGHT_WHITE = 107
  51. # SGR sequence pattern: CSI followed by params (digits, semicolons, colons) ending with 'm'
  52. # Colons are used in ITU T.416 (ISO 8613-6) extended color format: 38:2::R:G:B
  53. # This colon format is less common than semicolon (38;2;R;G;B) but supported by kitty,
  54. # iTerm2, and newer VTE-based terminals.
  55. _SGR_PATTERN = re.compile(r'\x1b\[([\d;:]*)m')
  56. # Fast path: quick check if any SGR sequence exists
  57. _SGR_QUICK_CHECK = re.compile(r'\x1b\[[\d;:]*m')
  58. # Reset sequence
  59. _SGR_RESET = '\x1b[0m'
  60. class _SGRState(NamedTuple):
  61. """
  62. Track active SGR terminal attributes by category (immutable).
  63. :param bold: Bold attribute (SGR 1).
  64. :param dim: Dim/faint attribute (SGR 2).
  65. :param italic: Italic attribute (SGR 3).
  66. :param underline: Underline attribute (SGR 4).
  67. :param blink: Slow blink attribute (SGR 5).
  68. :param rapid_blink: Rapid blink attribute (SGR 6).
  69. :param inverse: Inverse/reverse attribute (SGR 7).
  70. :param hidden: Hidden/invisible attribute (SGR 8).
  71. :param strikethrough: Strikethrough attribute (SGR 9).
  72. :param double_underline: Double underline attribute (SGR 21).
  73. :param foreground: Foreground color as tuple of SGR params, or None for default.
  74. :param background: Background color as tuple of SGR params, or None for default.
  75. """
  76. bold: bool = False
  77. dim: bool = False
  78. italic: bool = False
  79. underline: bool = False
  80. blink: bool = False
  81. rapid_blink: bool = False
  82. inverse: bool = False
  83. hidden: bool = False
  84. strikethrough: bool = False
  85. double_underline: bool = False
  86. foreground: tuple[int, ...] | None = None
  87. background: tuple[int, ...] | None = None
  88. # Default state with no attributes set
  89. _SGR_STATE_DEFAULT = _SGRState()
  90. def _sgr_state_is_active(state: _SGRState) -> bool:
  91. """
  92. Return True if any attributes are set.
  93. :param state: The SGR state to check.
  94. :returns: True if any attribute differs from default.
  95. """
  96. return (state.bold or state.dim or state.italic or state.underline
  97. or state.blink or state.rapid_blink or state.inverse or state.hidden
  98. or state.strikethrough or state.double_underline
  99. or state.foreground is not None or state.background is not None)
  100. def _sgr_state_to_sequence(state: _SGRState) -> str:
  101. """
  102. Generate minimal SGR sequence to restore this state from reset.
  103. :param state: The SGR state to convert.
  104. :returns: SGR escape sequence string, or empty string if no attributes set.
  105. """
  106. if not _sgr_state_is_active(state):
  107. return ''
  108. # Map boolean attributes to their SGR codes
  109. bool_attrs = [
  110. (state.bold, '1'), (state.dim, '2'), (state.italic, '3'),
  111. (state.underline, '4'), (state.blink, '5'), (state.rapid_blink, '6'),
  112. (state.inverse, '7'), (state.hidden, '8'), (state.strikethrough, '9'),
  113. (state.double_underline, '21'),
  114. ]
  115. params = [code for active, code in bool_attrs if active]
  116. # Add color params (already formatted as tuples)
  117. if state.foreground is not None:
  118. params.append(';'.join(str(p) for p in state.foreground))
  119. if state.background is not None:
  120. params.append(';'.join(str(p) for p in state.background))
  121. return f'\x1b[{";".join(params)}m'
  122. def _parse_sgr_params(sequence: str) -> list[int | tuple[int, ...]]:
  123. r"""
  124. Parse SGR sequence and return list of parameter values.
  125. Handles compound sequences like ``\x1b[1;31;4m`` -> [1, 31, 4].
  126. Empty params (e.g., ``\x1b[m``) are treated as [0] (reset).
  127. Colon-separated extended colors like ``\x1b[38:2::255:0:0m`` are returned
  128. as tuples: [(38, 2, 255, 0, 0)].
  129. :param sequence: SGR escape sequence string.
  130. :returns: List of integer parameters or tuples for colon-separated colors.
  131. """
  132. match = _SGR_PATTERN.match(sequence)
  133. if not match:
  134. return []
  135. params_str = match.group(1)
  136. if not params_str:
  137. return [0] # \x1b[m is equivalent to \x1b[0m
  138. result: list[int | tuple[int, ...]] = []
  139. for param in params_str.split(';'):
  140. if ':' in param:
  141. # Colon-separated extended color (ITU T.416 format)
  142. # e.g., "38:2::255:0:0" or "38:2:1:255:0:0" (with colorspace)
  143. parts = [int(p) if p else 0 for p in param.split(':')]
  144. result.append(tuple(parts))
  145. else:
  146. result.append(int(param) if param else 0)
  147. return result
  148. def _parse_extended_color(
  149. params: Iterator[int | tuple[int, ...]], base: int
  150. ) -> tuple[int, ...] | None:
  151. """
  152. Parse extended color (256-color or RGB) from parameter iterator.
  153. :param params: Iterator of remaining SGR parameters (semicolon-separated format).
  154. :param base: Base code (38 for foreground, 48 for background).
  155. :returns: Color tuple like (38, 5, N) or (38, 2, R, G, B), or None if malformed.
  156. """
  157. try:
  158. mode = next(params)
  159. if isinstance(mode, tuple):
  160. return None # Unexpected tuple, colon format handled separately
  161. if mode == 5: # 256-color
  162. n = next(params)
  163. if isinstance(n, tuple):
  164. return None
  165. return (int(base), 5, n)
  166. if mode == 2: # RGB
  167. r, g, b = next(params), next(params), next(params)
  168. if isinstance(r, tuple) or isinstance(g, tuple) or isinstance(b, tuple):
  169. return None
  170. return (int(base), 2, r, g, b)
  171. except StopIteration:
  172. pass
  173. return None
  174. def _sgr_state_update(state: _SGRState, sequence: str) -> _SGRState:
  175. # pylint: disable=too-many-branches,too-complex,too-many-statements
  176. # NOTE: When minimum Python version is 3.10+, this can be simplified using match/case.
  177. """
  178. Parse SGR sequence and return new state with updates applied.
  179. :param state: Current SGR state.
  180. :param sequence: SGR escape sequence string.
  181. :returns: New SGRState with updates applied.
  182. """
  183. params_list = _parse_sgr_params(sequence)
  184. params = iter(params_list)
  185. for p in params:
  186. # Handle colon-separated extended colors (ITU T.416 format)
  187. if isinstance(p, tuple):
  188. if len(p) >= 2 and p[0] == _SGR.FG_EXTENDED:
  189. # Foreground: (38, 2, [colorspace,] R, G, B) or (38, 5, N)
  190. state = state._replace(foreground=p)
  191. elif len(p) >= 2 and p[0] == _SGR.BG_EXTENDED:
  192. # Background: (48, 2, [colorspace,] R, G, B) or (48, 5, N)
  193. state = state._replace(background=p)
  194. continue
  195. if p == _SGR.RESET:
  196. state = _SGR_STATE_DEFAULT
  197. # Attribute ON codes
  198. elif p == _SGR.BOLD:
  199. state = state._replace(bold=True)
  200. elif p == _SGR.DIM:
  201. state = state._replace(dim=True)
  202. elif p == _SGR.ITALIC:
  203. state = state._replace(italic=True)
  204. elif p == _SGR.UNDERLINE:
  205. state = state._replace(underline=True)
  206. elif p == _SGR.BLINK:
  207. state = state._replace(blink=True)
  208. elif p == _SGR.RAPID_BLINK:
  209. state = state._replace(rapid_blink=True)
  210. elif p == _SGR.INVERSE:
  211. state = state._replace(inverse=True)
  212. elif p == _SGR.HIDDEN:
  213. state = state._replace(hidden=True)
  214. elif p == _SGR.STRIKETHROUGH:
  215. state = state._replace(strikethrough=True)
  216. elif p == _SGR.DOUBLE_UNDERLINE:
  217. state = state._replace(double_underline=True)
  218. # Attribute OFF codes
  219. elif p == _SGR.BOLD_DIM_OFF:
  220. state = state._replace(bold=False, dim=False)
  221. elif p == _SGR.ITALIC_OFF:
  222. state = state._replace(italic=False)
  223. elif p == _SGR.UNDERLINE_OFF:
  224. state = state._replace(underline=False, double_underline=False)
  225. elif p == _SGR.BLINK_OFF:
  226. state = state._replace(blink=False, rapid_blink=False)
  227. elif p == _SGR.INVERSE_OFF:
  228. state = state._replace(inverse=False)
  229. elif p == _SGR.HIDDEN_OFF:
  230. state = state._replace(hidden=False)
  231. elif p == _SGR.STRIKETHROUGH_OFF:
  232. state = state._replace(strikethrough=False)
  233. # Basic colors (30-37, 40-47 standard; 90-97, 100-107 bright)
  234. elif (_SGR.FG_BLACK <= p <= _SGR.FG_WHITE
  235. or _SGR.FG_BRIGHT_BLACK <= p <= _SGR.FG_BRIGHT_WHITE):
  236. state = state._replace(foreground=(p,))
  237. elif (_SGR.BG_BLACK <= p <= _SGR.BG_WHITE
  238. or _SGR.BG_BRIGHT_BLACK <= p <= _SGR.BG_BRIGHT_WHITE):
  239. state = state._replace(background=(p,))
  240. elif p == _SGR.FG_DEFAULT:
  241. state = state._replace(foreground=None)
  242. elif p == _SGR.BG_DEFAULT:
  243. state = state._replace(background=None)
  244. # Extended colors (semicolon-separated format)
  245. elif p == _SGR.FG_EXTENDED:
  246. if color := _parse_extended_color(params, _SGR.FG_EXTENDED):
  247. state = state._replace(foreground=color)
  248. elif p == _SGR.BG_EXTENDED:
  249. if color := _parse_extended_color(params, _SGR.BG_EXTENDED):
  250. state = state._replace(background=color)
  251. return state
  252. def propagate_sgr(lines: Sequence[str]) -> list[str]:
  253. r"""
  254. Propagate SGR codes across wrapped lines.
  255. When text with SGR styling is wrapped across multiple lines, each line
  256. needs to be self-contained for proper display. This function:
  257. - Ends each line with ``\x1b[0m`` if styles are active (prevents bleeding)
  258. - Starts each subsequent line with the active style restored
  259. :param lines: List of text lines, possibly containing SGR sequences.
  260. :returns: List of lines with SGR codes propagated.
  261. Example::
  262. >>> propagate_sgr(['\x1b[31mhello', 'world\x1b[0m'])
  263. ['\x1b[31mhello\x1b[0m', '\x1b[31mworld\x1b[0m']
  264. This is useful in cases of making special editors and viewers, and is used for the
  265. default modes (propagate_sgr=True) of :func:`wcwidth.width` and :func:`wcwidth.clip`.
  266. When wrapping and clipping text containing SGR sequences, maybe a previous line enabled the BLUE
  267. color--if we are viewing *only* the line following, we would want the carry over the BLUE color,
  268. and all lines with sequences should end with terminating reset (``\x1b[0m``).
  269. """
  270. # Fast path: check if any line contains SGR sequences
  271. if not any(_SGR_QUICK_CHECK.search(line) for line in lines) or not lines:
  272. return list(lines)
  273. result: list[str] = []
  274. state = _SGR_STATE_DEFAULT
  275. for line in lines:
  276. # Prefix with restoration sequence if state is active
  277. prefix = _sgr_state_to_sequence(state)
  278. # Update state by processing all SGR sequences in this line
  279. for match in _SGR_PATTERN.finditer(line):
  280. state = _sgr_state_update(state, match.group())
  281. # Build output line
  282. output_line = prefix + line if prefix else line
  283. if _sgr_state_is_active(state):
  284. output_line = output_line + _SGR_RESET
  285. result.append(output_line)
  286. return result