diff --git a/prompt_toolkit/formatted_text/ansi.py b/prompt_toolkit/formatted_text/ansi.py index 3d5706335..cebac9210 100644 --- a/prompt_toolkit/formatted_text/ansi.py +++ b/prompt_toolkit/formatted_text/ansi.py @@ -1,4 +1,5 @@ -from typing import Generator, List, Optional +from string import Formatter +from typing import Generator, List, Optional, Tuple, Union from prompt_toolkit.output.vt100 import BG_ANSI_COLORS, FG_ANSI_COLORS from prompt_toolkit.output.vt100 import _256_colors as _256_colors_table @@ -257,11 +258,17 @@ 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.) """ - # Escape all the arguments. - args = tuple(ansi_escape(a) for a in args) - kwargs = {k: ansi_escape(v) for k, v in kwargs.items()} + return ANSI(FORMATTER.vformat(self.value, args, kwargs)) - return ANSI(self.value.format(*args, **kwargs)) + def __mod__(self, value: object) -> "ANSI": + """ + ANSI('%s') % 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. @@ -275,8 +282,16 @@ def format(self, *args: str, **kwargs: str) -> "ANSI": _256_colors[i] = "#%02x%02x%02x" % (r, g, b) -def ansi_escape(text: str) -> str: +def ansi_escape(text: object) -> str: """ Replace characters with a special meaning. """ - return text.replace("\x1b", "?").replace("\b", "?") + 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() diff --git a/prompt_toolkit/formatted_text/html.py b/prompt_toolkit/formatted_text/html.py index 9a5213227..735ba2f1e 100644 --- a/prompt_toolkit/formatted_text/html.py +++ b/prompt_toolkit/formatted_text/html.py @@ -110,7 +110,7 @@ def format(self, *args: object, **kwargs: object) -> "HTML": """ return HTML(FORMATTER.vformat(self.value, args, kwargs)) - def __mod__(self, value: Union[object, Tuple[object, ...]]) -> "HTML": + def __mod__(self, value: object) -> "HTML": """ HTML('%s') % value """ diff --git a/tests/test_formatted_text.py b/tests/test_formatted_text.py index a49c67202..8b4924f96 100644 --- a/tests/test_formatted_text.py +++ b/tests/test_formatted_text.py @@ -103,6 +103,71 @@ def test_ansi_true_color(): ] +def test_ansi_interpolation(): + # %-style interpolation. + value = ANSI("\x1b[1m%s\x1b[0m") % "hello\x1b" + assert to_formatted_text(value) == [ + ("bold", "h"), + ("bold", "e"), + ("bold", "l"), + ("bold", "l"), + ("bold", "o"), + ("bold", "?"), + ] + + value = ANSI("\x1b[1m%s\x1b[0m") % ("\x1bhello",) + assert to_formatted_text(value) == [ + ("bold", "?"), + ("bold", "h"), + ("bold", "e"), + ("bold", "l"), + ("bold", "l"), + ("bold", "o"), + ] + + value = ANSI("\x1b[32m%s\x1b[45m%s") % ("He", "\x1bllo") + assert to_formatted_text(value) == [ + ("ansigreen", "H"), + ("ansigreen", "e"), + ("ansigreen bg:ansimagenta", "?"), + ("ansigreen bg:ansimagenta", "l"), + ("ansigreen bg:ansimagenta", "l"), + ("ansigreen bg:ansimagenta", "o"), + ] + + # Format function. + value = ANSI("\x1b[32m{0}\x1b[45m{1}").format("He\x1b", "llo") + assert to_formatted_text(value) == [ + ("ansigreen", "H"), + ("ansigreen", "e"), + ("ansigreen", "?"), + ("ansigreen bg:ansimagenta", "l"), + ("ansigreen bg:ansimagenta", "l"), + ("ansigreen bg:ansimagenta", "o"), + ] + + value = ANSI("\x1b[32m{a}\x1b[45m{b}").format(a="\x1bHe", b="llo") + assert to_formatted_text(value) == [ + ("ansigreen", "?"), + ("ansigreen", "H"), + ("ansigreen", "e"), + ("ansigreen bg:ansimagenta", "l"), + ("ansigreen bg:ansimagenta", "l"), + ("ansigreen bg:ansimagenta", "o"), + ] + + value = ANSI("\x1b[32m{:02d}\x1b[45m{:.3f}").format(3, 3.14159) + assert to_formatted_text(value) == [ + ("ansigreen", "0"), + ("ansigreen", "3"), + ("ansigreen bg:ansimagenta", "3"), + ("ansigreen bg:ansimagenta", "."), + ("ansigreen bg:ansimagenta", "1"), + ("ansigreen bg:ansimagenta", "4"), + ("ansigreen bg:ansimagenta", "2"), + ] + + def test_interpolation(): value = Template(" {} ").format(HTML("hello")) @@ -125,18 +190,18 @@ def test_interpolation(): def test_html_interpolation(): # %-style interpolation. - value = HTML("%s") % "hello" - assert to_formatted_text(value) == [("class:b", "hello")] + value = HTML("%s") % "&hello" + assert to_formatted_text(value) == [("class:b", "&hello")] - value = HTML("%s") % ("hello",) - assert to_formatted_text(value) == [("class:b", "hello")] + value = HTML("%s") % ("",) + assert to_formatted_text(value) == [("class:b", "")] - value = HTML("%s%s") % ("hello", "world") - assert to_formatted_text(value) == [("class:b", "hello"), ("class:u", "world")] + value = HTML("%s%s") % ("", "") + assert to_formatted_text(value) == [("class:b", ""), ("class:u", "")] # Format function. - value = HTML("{0}{1}").format("hello", "world") - assert to_formatted_text(value) == [("class:b", "hello"), ("class:u", "world")] + value = HTML("{0}{1}").format("'hello'", '"world"') + assert to_formatted_text(value) == [("class:b", "'hello'"), ("class:u", '"world"')] value = HTML("{a}{b}").format(a="hello", b="world") assert to_formatted_text(value) == [("class:b", "hello"), ("class:u", "world")]