diff --git a/CHANGELOG.md b/CHANGELOG.md index 47bf49b0..df615099 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/tests/test_write.py b/tests/test_write.py index 228d250d..e0590cb5 100644 --- a/tests/test_write.py +++ b/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" diff --git a/tomlkit/_utils.py b/tomlkit/_utils.py index d91c8309..1aa374f9 100644 --- a/tomlkit/_utils.py +++ b/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: diff --git a/tomlkit/items.py b/tomlkit/items.py index 51eba39e..4d56f564 100644 --- a/tomlkit/items.py +++ b/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