Speed up escaping strings and keys when writing TOML (#709)

Author: deadstrobe5

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Changed
 
 - Speed up repeated array appends by keeping existing indexes instead of rebuilding them.
+- Speed up creating string values and keys, e.g. when building a document from Python data and dumping it: strings are escaped in one pass instead of per character, and the escape and bare-key character sets are built once.
 
 ### Fixed
 

tests/test_write.py

@@ -53,6 +53,18 @@ def test_escape_special_characters_in_key() -> None:
     assert result == "baz"
 
 
+def test_write_escapes_every_sequence_in_a_string() -> None:
+    value = 'say "hi" \\ \x00\t\x7f """"" end'
+    d = {"a": value, "m": string(value + "\n", multiline=True)}
+    expected = (
+        'a = "say \\"hi\\" \\\\ \\u0000\\t\\u007f \\"\\"\\"\\"\\" end"\n'
+        'm = """say "hi" \\\\ \\u0000\\t\\u007f ""\\""" end\n"""\n'
+    )
+
+    assert expected == dumps(d)
+    assert loads(dumps(d)) == d
+
+
 def test_write_inline_table_in_nested_arrays() -> None:
     d = {"foo": [[{"a": 1}]]}
     expected = "foo = [[{a = 1}]]\n"

tomlkit/_utils.py

@@ -136,33 +136,25 @@ def _unicode_escape(seq: str) -> str:
     return "".join(f"\\u{ord(c):04x}" for c in seq)
 
 
-def escape_string(s: str, escape_sequences: Collection[str] = _basic_escapes) -> str:
-    s = decode(s)
-
-    res = []
-    start = 0
-
-    def flush(inc: int = 1) -> int:
-        if start != i:
-            res.append(s[start:i])
+def _escape_match(match: re.Match[str]) -> str:
+    seq = match.group()
+    return _compact_escapes.get(seq) or _unicode_escape(seq)
 
-        return i + inc
 
-    found_sequences = {seq for seq in escape_sequences if seq in s}
-
-    i = 0
-    while i < len(s):
-        for seq in found_sequences:
-            seq_len = len(seq)
-            if s[i:].startswith(seq):
-                start = flush(seq_len)
-                res.append(_compact_escapes.get(seq) or _unicode_escape(seq))
-                i += seq_len - 1  # fast-forward escape sequence
-        i += 1
-
-    flush()
+def escape_string(s: str, escape_sequences: Collection[str] = _basic_escapes) -> str:
+    s = decode(s)
 
-    return "".join(res)
+    found_sequences = [seq for seq in escape_sequences if seq in s]
+    if not found_sequences:
+        # always return a plain str, also for str subclasses such as String
+        return str(s)
+
+    # Replace every sequence in a single left-to-right pass. Longer sequences
+    # come first so that '"""' is matched as a whole.
+    pattern = "|".join(
+        re.escape(seq) for seq in sorted(found_sequences, key=len, reverse=True)
+    )
+    return re.sub(pattern, _escape_match, s)
 
 
 def merge_dicts(d1: dict[str, Any], d2: dict[str, Any]) -> None:

tomlkit/items.py

@@ -306,6 +306,8 @@ def _format_datetime(value: datetime) -> str:
 # a control character that a basic string must escape.
 _BARE_CR = re.compile(r"\r(?!\n)")
 
+_BARE_KEY_CHARS = frozenset(string.ascii_letters + string.digits + "-_")
+
 
 class StringType(Enum):
     # Single Line Basic
@@ -329,27 +331,14 @@ def select(cls, literal: bool = False, multiline: bool = False) -> StringType:
     @property
     def escaped_sequences(self) -> Collection[str]:
         # https://toml.io/en/v1.0.0#string
-        escaped_in_basic = CONTROL_CHARS | {"\\"}
-        allowed_in_multiline = {"\n", "\r"}
-        return {
-            StringType.SLB: escaped_in_basic | {'"'},
-            StringType.MLB: (escaped_in_basic | {'"""'}) - allowed_in_multiline,
-            StringType.SLL: (),
-            StringType.MLL: (),
-        }[self]
+        return _ESCAPED_SEQUENCES[self]
 
     @property
     def invalid_sequences(self) -> Collection[str]:
         # https://toml.io/en/v1.0.0#string
-        forbidden_in_literal = CONTROL_CHARS - {"\t"}
         # CR stays forbidden in multiline: the value is checked with CRLF
         # pairs collapsed, so any remaining CR is a lone control character
-        return {
-            StringType.SLB: (),
-            StringType.MLB: (),
-            StringType.SLL: forbidden_in_literal | {"'"},
-            StringType.MLL: (forbidden_in_literal | {"'''"}) - {"\n"},
-        }[self]
+        return _INVALID_SEQUENCES[self]
 
     @property
     def unit(self) -> str:
@@ -376,6 +365,23 @@ def toggle(self) -> StringType:
         }[self]
 
 
+_ESCAPED_IN_BASIC = CONTROL_CHARS | {"\\"}
+_FORBIDDEN_IN_LITERAL = CONTROL_CHARS - {"\t"}
+_ESCAPED_SEQUENCES: dict[StringType, Collection[str]] = {
+    StringType.SLB: _ESCAPED_IN_BASIC | {'"'},
+    # raw line breaks are allowed in multiline strings
+    StringType.MLB: (_ESCAPED_IN_BASIC | {'"""'}) - {"\n", "\r"},
+    StringType.SLL: (),
+    StringType.MLL: (),
+}
+_INVALID_SEQUENCES: dict[StringType, Collection[str]] = {
+    StringType.SLB: (),
+    StringType.MLB: (),
+    StringType.SLL: _FORBIDDEN_IN_LITERAL | {"'"},
+    StringType.MLL: (_FORBIDDEN_IN_LITERAL | {"'''"}) - {"\n"},
+}
+
+
 class BoolType(Enum):
     TRUE = "true"
     FALSE = "false"
@@ -480,9 +486,7 @@ def __init__(
             raise TypeError("Keys must be strings")
 
         if t is None:
-            if not k or any(
-                c not in string.ascii_letters + string.digits + "-" + "_" for c in k
-            ):
+            if not k or not _BARE_KEY_CHARS.issuperset(k):
                 t = KeyType.Basic
             else:
                 t = KeyType.Bare