summaryrefslogtreecommitdiffstats
path: root/src/prompt_toolkit/formatted_text/ansi.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/prompt_toolkit/formatted_text/ansi.py')
-rw-r--r--src/prompt_toolkit/formatted_text/ansi.py299
1 files changed, 299 insertions, 0 deletions
diff --git a/src/prompt_toolkit/formatted_text/ansi.py b/src/prompt_toolkit/formatted_text/ansi.py
new file mode 100644
index 0000000..08ec0b3
--- /dev/null
+++ b/src/prompt_toolkit/formatted_text/ansi.py
@@ -0,0 +1,299 @@
+from __future__ import annotations
+
+from string import Formatter
+from typing import Generator
+
+from prompt_toolkit.output.vt100 import BG_ANSI_COLORS, FG_ANSI_COLORS
+from prompt_toolkit.output.vt100 import _256_colors as _256_colors_table
+
+from .base import StyleAndTextTuples
+
+__all__ = [
+ "ANSI",
+ "ansi_escape",
+]
+
+
+class ANSI:
+ """
+ ANSI formatted text.
+ Take something ANSI escaped text, for use as a formatted string. E.g.
+
+ ::
+
+ ANSI('\\x1b[31mhello \\x1b[32mworld')
+
+ Characters between ``\\001`` and ``\\002`` are supposed to have a zero width
+ when printed, but these are literally sent to the terminal output. This can
+ be used for instance, for inserting Final Term prompt commands. They will
+ be translated into a prompt_toolkit '[ZeroWidthEscape]' fragment.
+ """
+
+ def __init__(self, value: str) -> None:
+ self.value = value
+ self._formatted_text: StyleAndTextTuples = []
+
+ # Default style attributes.
+ self._color: str | None = None
+ self._bgcolor: str | None = None
+ self._bold = False
+ self._underline = False
+ self._strike = False
+ self._italic = False
+ self._blink = False
+ self._reverse = False
+ self._hidden = False
+
+ # Process received text.
+ parser = self._parse_corot()
+ parser.send(None) # type: ignore
+ for c in value:
+ parser.send(c)
+
+ def _parse_corot(self) -> Generator[None, str, None]:
+ """
+ Coroutine that parses the ANSI escape sequences.
+ """
+ style = ""
+ formatted_text = self._formatted_text
+
+ while True:
+ # NOTE: CSI is a special token within a stream of characters that
+ # introduces an ANSI control sequence used to set the
+ # style attributes of the following characters.
+ csi = False
+
+ c = yield
+
+ # Everything between \001 and \002 should become a ZeroWidthEscape.
+ if c == "\001":
+ escaped_text = ""
+ while c != "\002":
+ c = yield
+ if c == "\002":
+ formatted_text.append(("[ZeroWidthEscape]", escaped_text))
+ c = yield
+ break
+ else:
+ escaped_text += c
+
+ # Check for CSI
+ if c == "\x1b":
+ # Start of color escape sequence.
+ square_bracket = yield
+ if square_bracket == "[":
+ csi = True
+ else:
+ continue
+ elif c == "\x9b":
+ csi = True
+
+ if csi:
+ # Got a CSI sequence. Color codes are following.
+ current = ""
+ params = []
+
+ while True:
+ char = yield
+
+ # Construct number
+ if char.isdigit():
+ current += char
+
+ # Eval number
+ else:
+ # Limit and save number value
+ params.append(min(int(current or 0), 9999))
+
+ # Get delimiter token if present
+ if char == ";":
+ current = ""
+
+ # Check and evaluate color codes
+ elif char == "m":
+ # Set attributes and token.
+ self._select_graphic_rendition(params)
+ style = self._create_style_string()
+ break
+
+ # Check and evaluate cursor forward
+ elif char == "C":
+ for i in range(params[0]):
+ # add <SPACE> using current style
+ formatted_text.append((style, " "))
+ break
+
+ else:
+ # Ignore unsupported sequence.
+ break
+ else:
+ # Add current character.
+ # NOTE: At this point, we could merge the current character
+ # into the previous tuple if the style did not change,
+ # however, it's not worth the effort given that it will
+ # be "Exploded" once again when it's rendered to the
+ # output.
+ formatted_text.append((style, c))
+
+ def _select_graphic_rendition(self, attrs: list[int]) -> None:
+ """
+ Taken a list of graphics attributes and apply changes.
+ """
+ if not attrs:
+ attrs = [0]
+ else:
+ attrs = list(attrs[::-1])
+
+ while attrs:
+ attr = attrs.pop()
+
+ if attr in _fg_colors:
+ self._color = _fg_colors[attr]
+ elif attr in _bg_colors:
+ self._bgcolor = _bg_colors[attr]
+ elif attr == 1:
+ self._bold = True
+ # elif attr == 2:
+ # self._faint = True
+ elif attr == 3:
+ self._italic = True
+ elif attr == 4:
+ self._underline = True
+ elif attr == 5:
+ self._blink = True # Slow blink
+ elif attr == 6:
+ self._blink = True # Fast blink
+ elif attr == 7:
+ self._reverse = True
+ elif attr == 8:
+ self._hidden = True
+ elif attr == 9:
+ self._strike = True
+ elif attr == 22:
+ self._bold = False # Normal intensity
+ elif attr == 23:
+ self._italic = False
+ elif attr == 24:
+ self._underline = False
+ elif attr == 25:
+ self._blink = False
+ elif attr == 27:
+ self._reverse = False
+ elif attr == 28:
+ self._hidden = False
+ elif attr == 29:
+ self._strike = False
+ elif not attr:
+ # Reset all style attributes
+ self._color = None
+ self._bgcolor = None
+ self._bold = False
+ self._underline = False
+ self._strike = False
+ self._italic = False
+ self._blink = False
+ self._reverse = False
+ self._hidden = False
+
+ elif attr in (38, 48) and len(attrs) > 1:
+ n = attrs.pop()
+
+ # 256 colors.
+ if n == 5 and len(attrs) >= 1:
+ if attr == 38:
+ m = attrs.pop()
+ self._color = _256_colors.get(m)
+ elif attr == 48:
+ m = attrs.pop()
+ self._bgcolor = _256_colors.get(m)
+
+ # True colors.
+ if n == 2 and len(attrs) >= 3:
+ try:
+ color_str = "#{:02x}{:02x}{:02x}".format(
+ attrs.pop(),
+ attrs.pop(),
+ attrs.pop(),
+ )
+ except IndexError:
+ pass
+ else:
+ if attr == 38:
+ self._color = color_str
+ elif attr == 48:
+ self._bgcolor = color_str
+
+ def _create_style_string(self) -> str:
+ """
+ Turn current style flags into a string for usage in a formatted text.
+ """
+ result = []
+ if self._color:
+ result.append(self._color)
+ if self._bgcolor:
+ result.append("bg:" + self._bgcolor)
+ if self._bold:
+ result.append("bold")
+ if self._underline:
+ result.append("underline")
+ if self._strike:
+ result.append("strike")
+ if self._italic:
+ result.append("italic")
+ if self._blink:
+ result.append("blink")
+ if self._reverse:
+ result.append("reverse")
+ if self._hidden:
+ result.append("hidden")
+
+ return " ".join(result)
+
+ def __repr__(self) -> str:
+ return f"ANSI({self.value!r})"
+
+ def __pt_formatted_text__(self) -> StyleAndTextTuples:
+ return self._formatted_text
+
+ def format(self, *args: str, **kwargs: str) -> ANSI:
+ """
+ Like `str.format`, but make sure that the arguments are properly
+ escaped. (No ANSI escapes can be injected.)
+ """
+ return ANSI(FORMATTER.vformat(self.value, args, kwargs))
+
+ def __mod__(self, value: object) -> ANSI:
+ """
+ ANSI('<b>%s</b>') % value
+ """
+ if not isinstance(value, tuple):
+ value = (value,)
+
+ value = tuple(ansi_escape(i) for i in value)
+ return ANSI(self.value % value)
+
+
+# Mapping of the ANSI color codes to their names.
+_fg_colors = {v: k for k, v in FG_ANSI_COLORS.items()}
+_bg_colors = {v: k for k, v in BG_ANSI_COLORS.items()}
+
+# Mapping of the escape codes for 256colors to their 'ffffff' value.
+_256_colors = {}
+
+for i, (r, g, b) in enumerate(_256_colors_table.colors):
+ _256_colors[i] = f"#{r:02x}{g:02x}{b:02x}"
+
+
+def ansi_escape(text: object) -> str:
+ """
+ Replace characters with a special meaning.
+ """
+ return str(text).replace("\x1b", "?").replace("\b", "?")
+
+
+class ANSIFormatter(Formatter):
+ def format_field(self, value: object, format_spec: str) -> str:
+ return ansi_escape(format(value, format_spec))
+
+
+FORMATTER = ANSIFormatter()