From f0dbaed5c97c3cb9f390df33cf4592eec77f0850 Mon Sep 17 00:00:00 2001 From: Dickson Date: Fri, 6 Nov 2026 15:10:40 -0500 Subject: [PATCH 1/2] Speed up key lookups and replacements in tables Replacing an existing key scanned every key in the container to recover the stored Key object, so updating many keys of a large table was quadratic. Take the stored key from the body slot that _map already points to instead. Looking up a plain str key built and validated a temporary SingleKey just to probe _map. SingleKey hashes and compares like its string, so probe _map with the string directly. Containers that have ever held a custom Key implementation keep the previous normalization and comparison scan, since such keys may not accept a str operand. Missing key errors still report the normalized key. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 1 + tests/test_container_updates.py | 123 ++++++++++++++++++++++++++++++++ tomlkit/container.py | 47 ++++++++---- 3 files changed, 158 insertions(+), 13 deletions(-) create mode 100644 tests/test_container_updates.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 47bf49b0..61e2ea63 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Changed +- Speed up looking up and replacing keys in tables: string lookups no longer build a temporary key, and replacing an existing value no longer scans every key in the table, which made updating many keys of a large table quadratic. - Speed up repeated array appends by keeping existing indexes instead of rebuilding them. ### Fixed diff --git a/tests/test_container_updates.py b/tests/test_container_updates.py new file mode 100644 index 00000000..75d8a592 --- /dev/null +++ b/tests/test_container_updates.py @@ -0,0 +1,123 @@ +import copy +import pickle + +from collections.abc import Callable +from typing import Any + +import pytest + +from tomlkit import key +from tomlkit import parse +from tomlkit.items import Item +from tomlkit.items import Key +from tomlkit.items import SingleKey +from tomlkit.toml_document import TOMLDocument + + +class StringKey(str): + pass + + +class KeyOperandOnly(SingleKey): + __hash__ = SingleKey.__hash__ + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, Key): + raise TypeError("Key operand required") + return super().__eq__(other) + + +def pickle_copy(doc: TOMLDocument) -> TOMLDocument: + return pickle.loads(pickle.dumps(doc)) # type: ignore[no-any-return] + + +@pytest.mark.parametrize("clone", [lambda x: x, copy.copy, copy.deepcopy, pickle_copy]) +@pytest.mark.parametrize("lookup", ["custom", "ordinary"]) +def test_custom_key_comparisons_preserved( + clone: Callable[[TOMLDocument], TOMLDocument], lookup: str +) -> None: + doc = parse("") + doc.append(KeyOperandOnly("custom"), 1) + doc.append("ordinary", 2) + doc = clone(doc) + assert "custom" in doc + assert doc["custom"] == 1 + assert doc["ordinary"] == 2 + # Upstream's assignment scan compares every stored key to the original + # query, including unrelated custom keys, before replacing anything. + with pytest.raises(TypeError, match="Key operand required"): + doc[lookup] = 3 + doc[key(lookup)] = 4 + assert doc[lookup] == 4 + + +@pytest.mark.parametrize("method", ["insert", "replace"]) +def test_custom_key_tracking_after_insert_or_replace(method: str) -> None: + doc = parse("first = 1\nlast = 2\n") + custom = KeyOperandOnly("custom") + if method == "insert": + doc._insert_at(0, custom, 3) + else: + value = doc.item("first") + assert isinstance(value, Item) + doc._replace("first", custom, value) + assert "custom" in doc + assert doc["custom"] in (1, 3) + with pytest.raises(TypeError, match="Key operand required"): + doc["custom"] = 4 + + +@pytest.mark.parametrize("missing", ["", "a b", 'a"b', "雪", "a.b", "\0"]) +def test_lookup_missing_key_preserves_quoted_error(missing: str) -> None: + doc = parse("present = 1\n") + with pytest.raises(KeyError) as bare_error: + doc[missing] + with pytest.raises(KeyError) as key_error: + doc[key(missing)] + assert type(bare_error.value) is type(key_error.value) + assert str(bare_error.value) == str(key_error.value) + assert missing not in doc + + +@pytest.mark.parametrize("lookup", ["a key", StringKey("a key"), key("a key")]) +def test_repeated_updates_preserve_canonical_key_and_body_order(lookup: Any) -> None: + doc = parse("first = 0\n 'a key' = 'old' # keep\nlast = 2\n") + original_key = doc.body[1][0] + for value in ("one", "two", "three"): + doc[lookup] = value + assert list(doc) == ["first", "a key", "last"] + assert doc.unwrap() == {"first": 0, "a key": value, "last": 2} + if isinstance(lookup, str): + assert doc.body[1][0] is original_key + assert doc.as_string() == ( + f"first = 0\n 'a key' = '{value}' # keep\nlast = 2\n" + ) + + +def test_update_after_deleted_slots_and_out_of_order_fragments() -> None: + doc = parse("[a.x]\nx = 1\n[b]\nb = 2\n[a.y]\ny = 3\n") + del doc["b"] + doc["a"] = {"replacement": 4} + assert list(doc) == ["a"] + assert doc.unwrap() == {"a": {"replacement": 4}} + assert parse(doc.as_string()).unwrap() == doc.unwrap() + doc["a"] = 5 + assert doc.unwrap() == {"a": 5} + assert parse(doc.as_string()).unwrap() == doc.unwrap() + + +def test_existing_key_update_does_not_compare_unrelated_keys( + monkeypatch: pytest.MonkeyPatch, +) -> None: + doc = parse("".join(f"dependency_{i} = 1\n" for i in range(100))) + comparisons = [] + original = SingleKey.__eq__ + + def compare(self: SingleKey, other: Any) -> bool: + comparisons.append(self.key) + return original(self, other) + + monkeypatch.setattr(SingleKey, "__eq__", compare) + doc["dependency_99"] = 2 + assert set(comparisons) == {"dependency_99"} + assert doc["dependency_99"] == 2 diff --git a/tomlkit/container.py b/tomlkit/container.py index 95ef73b2..b287b5c8 100644 --- a/tomlkit/container.py +++ b/tomlkit/container.py @@ -111,6 +111,10 @@ class Container(_CustomDict): # type: ignore[type-arg] def __init__(self, parsed: bool = False) -> None: self._map: dict[Key, int | tuple[int, ...]] = {} + # SingleKey hashes and compares like its string value. Custom Key + # implementations may require a Key operand, so keep their lookups on + # the normalization path, even after such a key has been removed. + self._has_custom_keys = False self._body: list[tuple[Key | None, Item]] = [] self._parsed = parsed self._table_keys: list[Key] = [] @@ -528,6 +532,8 @@ def _validate_table_candidate(self, current: Table, candidate: Table) -> None: raise TOMLKitError("Redefinition of an existing table") def _raw_append(self, key: Key | None, item: Item) -> None: + if key is not None and type(key) is not SingleKey: + self._has_custom_keys = True if key is not None and key in self._map: current_idx = self._map[key] if not isinstance(current_idx, tuple): @@ -596,6 +602,9 @@ def _insert_at(self, idx: int, key: Key | str, item: Any) -> Container: if not isinstance(key, Key): key = SingleKey(key) + if type(key) is not SingleKey: + self._has_custom_keys = True + item = _item(item) if idx > 0: @@ -638,11 +647,13 @@ def _insert_at(self, idx: int, key: Key | str, item: Any) -> Container: def item(self, key: Key | str) -> Item | OutOfOrderTableProxy: """Get an item for the given key.""" - if not isinstance(key, Key): + if not isinstance(key, Key) and (type(key) is not str or self._has_custom_keys): key = SingleKey(key) - idx = self._map.get(key) + idx = self._map.get(key) # type: ignore[arg-type] if idx is None: + if isinstance(key, str): + key = SingleKey(key) raise NonExistentKey(key) if isinstance(idx, tuple): @@ -849,16 +860,12 @@ def __getitem__(self, key: Key | str) -> Any: return item def __contains__(self, key: object) -> bool: - # Native membership test. The inherited ``MutableMapping.__contains__`` - # resolves the value via ``__getitem__``/``item()`` (and builds a - # ``NonExistentKey`` on every absent key) only to discard it. Resolve the - # key the same way ``item()`` does -- ``str`` becomes a ``SingleKey`` - # (a non-str/non-``Key`` argument still raises ``TypeError``) -- then - # probe ``_map`` directly. For an out-of-order table the proxy is still - # built so its validation runs exactly as before. - if not isinstance(key, Key): + # Exact strings can probe standard SingleKeys directly. Other queries + # retain normalization, including its TypeError for invalid key types. + # Out-of-order tables must still validate their fragments. + if not isinstance(key, Key) and (type(key) is not str or self._has_custom_keys): key = SingleKey(key) # type: ignore[arg-type] - idx = self._map.get(key) + idx = self._map.get(key) # type: ignore[arg-type] if idx is None: return False if isinstance(idx, tuple): @@ -866,7 +873,7 @@ def __contains__(self, key: object) -> bool: # if the validation cache covers exactly these fragments they # are already known to merge cleanly, so skip rebuilding the # proxy (which is linear in the number of fragments) - validated, _ = self._validation_cache.get(key, (0, None)) + validated, _ = self._validation_cache.get(key, (0, None)) # type: ignore[arg-type] if not (self._parsed and validated == len(idx)): OutOfOrderTableProxy(self, idx) return True @@ -888,7 +895,17 @@ def _item_for_key( def __setitem__(self, key: Key | str, value: Any) -> None: if key in self: - old_key = next(filter(lambda k: k == key, self._map)) + old_key: Key | None + if self._has_custom_keys or ( + type(key) is not str and type(key) is not SingleKey + ): + old_key = next(filter(lambda k: k == key, self._map)) + else: + idx = self._map[key] # type: ignore[index] + if isinstance(idx, tuple): + idx = idx[0] + old_key = self._body[idx][0] + assert old_key is not None value = self._item_for_key(key, value) self._replace(old_key, key, value) else: @@ -957,6 +974,8 @@ def _replace_at( else: # Inherit the sep of the old key new_key = k + if type(new_key) is not SingleKey: + self._has_custom_keys = True del self._map[k] self._map[new_key] = idx if new_key != k: @@ -1050,6 +1069,7 @@ def __reduce_ex__(self, protocol: int) -> tuple[type, tuple[bool], tuple[Any, .. def __setstate__(self, state: tuple[Any, ...]) -> None: self._map = state[0] + self._has_custom_keys = any(type(k) is not SingleKey for k in self._map) self._body = state[1] self._parsed = state[2] self._table_keys = state[3] @@ -1071,6 +1091,7 @@ def __copy__(self) -> Self: c._body += self.body c._map.update(self._map) + c._has_custom_keys = self._has_custom_keys c._out_of_order_keys |= self._out_of_order_keys return c From 05a82462e8bae7d6202caf43019412635bee62db Mon Sep 17 00:00:00 2001 From: Dickson Date: Thu, 12 Nov 2026 22:13:18 -0500 Subject: [PATCH 2/2] Drop the fallback for custom Key implementations A container only stores SingleKeys (dotted keys are split before they are stored), so plain str queries can always probe _map directly, and replacing a key can always take the stored key from its body slot. This removes the _has_custom_keys flag and the comparison scan it guarded. Co-Authored-By: Claude Opus 5.5 --- tests/test_container_updates.py | 56 --------------------------------- tomlkit/container.py | 40 +++++++---------------- 2 files changed, 12 insertions(+), 84 deletions(-) diff --git a/tests/test_container_updates.py b/tests/test_container_updates.py index 75d8a592..620cca85 100644 --- a/tests/test_container_updates.py +++ b/tests/test_container_updates.py @@ -1,72 +1,16 @@ -import copy -import pickle - -from collections.abc import Callable from typing import Any import pytest from tomlkit import key from tomlkit import parse -from tomlkit.items import Item -from tomlkit.items import Key from tomlkit.items import SingleKey -from tomlkit.toml_document import TOMLDocument class StringKey(str): pass -class KeyOperandOnly(SingleKey): - __hash__ = SingleKey.__hash__ - - def __eq__(self, other: Any) -> bool: - if not isinstance(other, Key): - raise TypeError("Key operand required") - return super().__eq__(other) - - -def pickle_copy(doc: TOMLDocument) -> TOMLDocument: - return pickle.loads(pickle.dumps(doc)) # type: ignore[no-any-return] - - -@pytest.mark.parametrize("clone", [lambda x: x, copy.copy, copy.deepcopy, pickle_copy]) -@pytest.mark.parametrize("lookup", ["custom", "ordinary"]) -def test_custom_key_comparisons_preserved( - clone: Callable[[TOMLDocument], TOMLDocument], lookup: str -) -> None: - doc = parse("") - doc.append(KeyOperandOnly("custom"), 1) - doc.append("ordinary", 2) - doc = clone(doc) - assert "custom" in doc - assert doc["custom"] == 1 - assert doc["ordinary"] == 2 - # Upstream's assignment scan compares every stored key to the original - # query, including unrelated custom keys, before replacing anything. - with pytest.raises(TypeError, match="Key operand required"): - doc[lookup] = 3 - doc[key(lookup)] = 4 - assert doc[lookup] == 4 - - -@pytest.mark.parametrize("method", ["insert", "replace"]) -def test_custom_key_tracking_after_insert_or_replace(method: str) -> None: - doc = parse("first = 1\nlast = 2\n") - custom = KeyOperandOnly("custom") - if method == "insert": - doc._insert_at(0, custom, 3) - else: - value = doc.item("first") - assert isinstance(value, Item) - doc._replace("first", custom, value) - assert "custom" in doc - assert doc["custom"] in (1, 3) - with pytest.raises(TypeError, match="Key operand required"): - doc["custom"] = 4 - - @pytest.mark.parametrize("missing", ["", "a b", 'a"b', "雪", "a.b", "\0"]) def test_lookup_missing_key_preserves_quoted_error(missing: str) -> None: doc = parse("present = 1\n") diff --git a/tomlkit/container.py b/tomlkit/container.py index b287b5c8..2c831d87 100644 --- a/tomlkit/container.py +++ b/tomlkit/container.py @@ -111,10 +111,6 @@ class Container(_CustomDict): # type: ignore[type-arg] def __init__(self, parsed: bool = False) -> None: self._map: dict[Key, int | tuple[int, ...]] = {} - # SingleKey hashes and compares like its string value. Custom Key - # implementations may require a Key operand, so keep their lookups on - # the normalization path, even after such a key has been removed. - self._has_custom_keys = False self._body: list[tuple[Key | None, Item]] = [] self._parsed = parsed self._table_keys: list[Key] = [] @@ -532,8 +528,6 @@ def _validate_table_candidate(self, current: Table, candidate: Table) -> None: raise TOMLKitError("Redefinition of an existing table") def _raw_append(self, key: Key | None, item: Item) -> None: - if key is not None and type(key) is not SingleKey: - self._has_custom_keys = True if key is not None and key in self._map: current_idx = self._map[key] if not isinstance(current_idx, tuple): @@ -602,9 +596,6 @@ def _insert_at(self, idx: int, key: Key | str, item: Any) -> Container: if not isinstance(key, Key): key = SingleKey(key) - if type(key) is not SingleKey: - self._has_custom_keys = True - item = _item(item) if idx > 0: @@ -647,7 +638,9 @@ def _insert_at(self, idx: int, key: Key | str, item: Any) -> Container: def item(self, key: Key | str) -> Item | OutOfOrderTableProxy: """Get an item for the given key.""" - if not isinstance(key, Key) and (type(key) is not str or self._has_custom_keys): + # SingleKey hashes and compares like its string, so a plain str can + # probe _map directly. + if type(key) is not str and not isinstance(key, Key): key = SingleKey(key) idx = self._map.get(key) # type: ignore[arg-type] @@ -860,10 +853,10 @@ def __getitem__(self, key: Key | str) -> Any: return item def __contains__(self, key: object) -> bool: - # Exact strings can probe standard SingleKeys directly. Other queries - # retain normalization, including its TypeError for invalid key types. + # Exact strings can probe _map directly. Other queries keep + # normalization, including its TypeError for invalid key types. # Out-of-order tables must still validate their fragments. - if not isinstance(key, Key) and (type(key) is not str or self._has_custom_keys): + if type(key) is not str and not isinstance(key, Key): key = SingleKey(key) # type: ignore[arg-type] idx = self._map.get(key) # type: ignore[arg-type] if idx is None: @@ -895,17 +888,12 @@ def _item_for_key( def __setitem__(self, key: Key | str, value: Any) -> None: if key in self: - old_key: Key | None - if self._has_custom_keys or ( - type(key) is not str and type(key) is not SingleKey - ): - old_key = next(filter(lambda k: k == key, self._map)) - else: - idx = self._map[key] # type: ignore[index] - if isinstance(idx, tuple): - idx = idx[0] - old_key = self._body[idx][0] - assert old_key is not None + # The body slot that _map points to holds the stored key. + idx = self._map[key] # type: ignore[index] + if isinstance(idx, tuple): + idx = idx[0] + old_key = self._body[idx][0] + assert old_key is not None value = self._item_for_key(key, value) self._replace(old_key, key, value) else: @@ -974,8 +962,6 @@ def _replace_at( else: # Inherit the sep of the old key new_key = k - if type(new_key) is not SingleKey: - self._has_custom_keys = True del self._map[k] self._map[new_key] = idx if new_key != k: @@ -1069,7 +1055,6 @@ def __reduce_ex__(self, protocol: int) -> tuple[type, tuple[bool], tuple[Any, .. def __setstate__(self, state: tuple[Any, ...]) -> None: self._map = state[0] - self._has_custom_keys = any(type(k) is not SingleKey for k in self._map) self._body = state[1] self._parsed = state[2] self._table_keys = state[3] @@ -1091,7 +1076,6 @@ def __copy__(self) -> Self: c._body += self.body c._map.update(self._map) - c._has_custom_keys = self._has_custom_keys c._out_of_order_keys |= self._out_of_order_keys return c