perf: fast-path sanitize_for_rich to skip per-call allocation

src/smolagents/utils.py

@@ -61,6 +61,11 @@ def _is_package_available(package_name: str) -> bool:
 ]
 
 
+# ASCII control characters that must be made visible: 0x00-0x1f and 0x7f, except the
+# common whitespace chars \t (0x09), \n (0x0a) and \r (0x0d) which are passed through.
+_CONTROL_CHAR_RE = re.compile("[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
+
+
 def sanitize_for_rich(value) -> str:
     """
     Convert arbitrary values (including bytes / control characters) into a safe string for Rich.
@@ -77,6 +82,15 @@ def sanitize_for_rich(value) -> str:
     else:
         s = str(value)
 
+    # Fast path: the common case is text with no control characters to escape. Skip the
+    # per-character list build + join. str() normalizes str subclasses back to a plain
+    # str -- a no-op returning the same object for an exact str (so the allocation is
+    # still avoided), while matching the slow path's "".join(...), which also always
+    # yields a plain str. This keeps callers like Rich from ever receiving a str subclass
+    # whose overridden methods (e.g. translate) could behave differently.
+    if not _CONTROL_CHAR_RE.search(s):
+        return str(s)
+
     out: list[str] = []
     for ch in s:
         code = ord(ch)

tests/test_utils.py

@@ -29,6 +29,7 @@
     is_valid_name,
     parse_code_blobs,
     parse_json_blob,
+    sanitize_for_rich,
 )
 
 
@@ -552,3 +553,94 @@ def test_agent_gradio_app_template_excludes_class_keyword():
         ast.parse(result)
     except SyntaxError as e:
         pytest.fail(f"Generated app.py contains syntax error: {e}")
+
+
+# ---------------------------------------------------------------------------
+# sanitize_for_rich
+# ---------------------------------------------------------------------------
+def _reference_sanitize(value):
+    """The pre-fast-path implementation: always builds a list and joins it, which
+    always yields a plain ``str``. Used as the equivalence oracle in the tests below."""
+    if value is None:
+        s = ""
+    elif isinstance(value, str):
+        s = value
+    elif isinstance(value, (bytes, bytearray, memoryview)):
+        s = bytes(value).decode("utf-8", errors="replace")
+    else:
+        s = str(value)
+    out = []
+    for ch in s:
+        code = ord(ch)
+        if ch in ("\n", "\t", "\r"):
+            out.append(ch)
+        elif code < 32 or code == 127:
+            out.append(f"\\x{code:02x}")
+        else:
+            out.append(ch)
+    return "".join(out)
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        None,
+        "",
+        "hello world",
+        "multi\nline\ttext\r\n",
+        "unicode: café — ✓ 日本語",
+        "brackets [bold]x[/bold]",
+        b"raw bytes \x00\x01\x1f\x7f end",
+        bytearray(b"\x02ab"),
+        "ctrl \x00\x07\x1b\x7f mixed",
+        "vtab\x0b formfeed\x0c",
+        12345,
+        3.14,
+        "\x00" * 5,
+        "a" * 200,
+    ],
+)
+def test_sanitize_for_rich_matches_reference(value):
+    """Output must be byte-identical to the list-build-and-join reference."""
+    assert sanitize_for_rich(value) == _reference_sanitize(value)
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        None,
+        "",
+        "plain control-free text",
+        "brackets [bold]x[/bold]",
+        "ctrl \x07 needs escaping",
+        b"bytes in",
+        12345,
+    ],
+)
+def test_sanitize_for_rich_always_returns_plain_str(value):
+    """Regression: the return value must always be an *exact* ``str`` (not a subclass).
+
+    The fast path returns ``str(s)`` rather than ``s`` so that a ``str`` subclass
+    passed in as arbitrary tool-log payload is normalized exactly like the slow
+    path's ``"".join(...)`` would. Callers hand the result to Rich ``Text(...)``,
+    and a subclass with overridden methods (e.g. ``translate``) could otherwise
+    behave differently from a plain ``str``.
+    """
+    result = sanitize_for_rich(value)
+    assert type(result) is str
+
+
+def test_sanitize_for_rich_normalizes_str_subclass():
+    """A control-free ``str`` subclass must come back as a plain ``str`` with equal
+    content, so its overridden methods can never reach Rich."""
+
+    class Weird(str):
+        def translate(self, *args, **kwargs):
+            raise ValueError("subclass translate must never be invoked by Rich")
+
+    payload = Weird("control-free subclass payload")
+    result = sanitize_for_rich(payload)
+    assert type(result) is str
+    assert result == "control-free subclass payload"
+    # The plain-str result routes through the built-in translate, never the override.
+    assert result.translate({0x07: None}) == "control-free subclass payload"