ansi.py 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240
  1. import re
  2. import sys
  3. from contextlib import suppress
  4. from typing import Iterable, NamedTuple, Optional
  5. from .color import Color
  6. from .style import Style
  7. from .text import Text
  8. re_ansi = re.compile(
  9. r"""
  10. (?:\x1b\](.*?)\x1b\\)|
  11. (?:\x1b([(@-Z\\-_]|\[[0-?]*[ -/]*[@-~]))
  12. """,
  13. re.VERBOSE,
  14. )
  15. class _AnsiToken(NamedTuple):
  16. """Result of ansi tokenized string."""
  17. plain: str = ""
  18. sgr: Optional[str] = ""
  19. osc: Optional[str] = ""
  20. def _ansi_tokenize(ansi_text: str) -> Iterable[_AnsiToken]:
  21. """Tokenize a string in to plain text and ANSI codes.
  22. Args:
  23. ansi_text (str): A String containing ANSI codes.
  24. Yields:
  25. AnsiToken: A named tuple of (plain, sgr, osc)
  26. """
  27. position = 0
  28. sgr: Optional[str]
  29. osc: Optional[str]
  30. for match in re_ansi.finditer(ansi_text):
  31. start, end = match.span(0)
  32. osc, sgr = match.groups()
  33. if start > position:
  34. yield _AnsiToken(ansi_text[position:start])
  35. if sgr:
  36. if sgr == "(":
  37. position = end + 1
  38. continue
  39. if sgr.endswith("m"):
  40. yield _AnsiToken("", sgr[1:-1], osc)
  41. else:
  42. yield _AnsiToken("", sgr, osc)
  43. position = end
  44. if position < len(ansi_text):
  45. yield _AnsiToken(ansi_text[position:])
  46. SGR_STYLE_MAP = {
  47. 1: "bold",
  48. 2: "dim",
  49. 3: "italic",
  50. 4: "underline",
  51. 5: "blink",
  52. 6: "blink2",
  53. 7: "reverse",
  54. 8: "conceal",
  55. 9: "strike",
  56. 21: "underline2",
  57. 22: "not dim not bold",
  58. 23: "not italic",
  59. 24: "not underline",
  60. 25: "not blink",
  61. 26: "not blink2",
  62. 27: "not reverse",
  63. 28: "not conceal",
  64. 29: "not strike",
  65. 30: "color(0)",
  66. 31: "color(1)",
  67. 32: "color(2)",
  68. 33: "color(3)",
  69. 34: "color(4)",
  70. 35: "color(5)",
  71. 36: "color(6)",
  72. 37: "color(7)",
  73. 39: "default",
  74. 40: "on color(0)",
  75. 41: "on color(1)",
  76. 42: "on color(2)",
  77. 43: "on color(3)",
  78. 44: "on color(4)",
  79. 45: "on color(5)",
  80. 46: "on color(6)",
  81. 47: "on color(7)",
  82. 49: "on default",
  83. 51: "frame",
  84. 52: "encircle",
  85. 53: "overline",
  86. 54: "not frame not encircle",
  87. 55: "not overline",
  88. 90: "color(8)",
  89. 91: "color(9)",
  90. 92: "color(10)",
  91. 93: "color(11)",
  92. 94: "color(12)",
  93. 95: "color(13)",
  94. 96: "color(14)",
  95. 97: "color(15)",
  96. 100: "on color(8)",
  97. 101: "on color(9)",
  98. 102: "on color(10)",
  99. 103: "on color(11)",
  100. 104: "on color(12)",
  101. 105: "on color(13)",
  102. 106: "on color(14)",
  103. 107: "on color(15)",
  104. }
  105. class AnsiDecoder:
  106. """Translate ANSI code in to styled Text."""
  107. def __init__(self) -> None:
  108. self.style = Style.null()
  109. def decode(self, terminal_text: str) -> Iterable[Text]:
  110. """Decode ANSI codes in an iterable of lines.
  111. Args:
  112. lines (Iterable[str]): An iterable of lines of terminal output.
  113. Yields:
  114. Text: Marked up Text.
  115. """
  116. for line in terminal_text.splitlines():
  117. yield self.decode_line(line)
  118. def decode_line(self, line: str) -> Text:
  119. """Decode a line containing ansi codes.
  120. Args:
  121. line (str): A line of terminal output.
  122. Returns:
  123. Text: A Text instance marked up according to ansi codes.
  124. """
  125. from_ansi = Color.from_ansi
  126. from_rgb = Color.from_rgb
  127. _Style = Style
  128. text = Text()
  129. append = text.append
  130. line = line.rsplit("\r", 1)[-1]
  131. for plain_text, sgr, osc in _ansi_tokenize(line):
  132. if plain_text:
  133. append(plain_text, self.style or None)
  134. elif osc is not None:
  135. if osc.startswith("8;"):
  136. _params, semicolon, link = osc[2:].partition(";")
  137. if semicolon:
  138. self.style = self.style.update_link(link or None)
  139. elif sgr is not None:
  140. # Translate in to semi-colon separated codes
  141. # Ignore invalid codes, because we want to be lenient
  142. codes = [
  143. min(255, int(_code) if _code else 0)
  144. for _code in sgr.split(";")
  145. if _code.isdigit() or _code == ""
  146. ]
  147. iter_codes = iter(codes)
  148. for code in iter_codes:
  149. if code == 0:
  150. # reset
  151. self.style = _Style.null()
  152. elif code in SGR_STYLE_MAP:
  153. # styles
  154. self.style += _Style.parse(SGR_STYLE_MAP[code])
  155. elif code == 38:
  156. #  Foreground
  157. with suppress(StopIteration):
  158. color_type = next(iter_codes)
  159. if color_type == 5:
  160. self.style += _Style.from_color(
  161. from_ansi(next(iter_codes))
  162. )
  163. elif color_type == 2:
  164. self.style += _Style.from_color(
  165. from_rgb(
  166. next(iter_codes),
  167. next(iter_codes),
  168. next(iter_codes),
  169. )
  170. )
  171. elif code == 48:
  172. # Background
  173. with suppress(StopIteration):
  174. color_type = next(iter_codes)
  175. if color_type == 5:
  176. self.style += _Style.from_color(
  177. None, from_ansi(next(iter_codes))
  178. )
  179. elif color_type == 2:
  180. self.style += _Style.from_color(
  181. None,
  182. from_rgb(
  183. next(iter_codes),
  184. next(iter_codes),
  185. next(iter_codes),
  186. ),
  187. )
  188. return text
  189. if sys.platform != "win32" and __name__ == "__main__": # pragma: no cover
  190. import io
  191. import os
  192. import pty
  193. import sys
  194. decoder = AnsiDecoder()
  195. stdout = io.BytesIO()
  196. def read(fd: int) -> bytes:
  197. data = os.read(fd, 1024)
  198. stdout.write(data)
  199. return data
  200. pty.spawn(sys.argv[1:], read)
  201. from .console import Console
  202. console = Console(record=True)
  203. stdout_result = stdout.getvalue().decode("utf-8")
  204. print(stdout_result)
  205. for line in decoder.decode(stdout_result):
  206. console.print(line)
  207. console.save_html("stdout.html")