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..620cca85 --- /dev/null +++ b/tests/test_container_updates.py @@ -0,0 +1,67 @@ +from typing import Any + +import pytest + +from tomlkit import key +from tomlkit import parse +from tomlkit.items import SingleKey + + +class StringKey(str): + pass + + +@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..2c831d87 100644 --- a/tomlkit/container.py +++ b/tomlkit/container.py @@ -638,11 +638,15 @@ 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): + # 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) + 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 +853,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 _map directly. Other queries keep + # normalization, including its TypeError for invalid key types. + # Out-of-order tables must still validate their fragments. + if type(key) is not str and not isinstance(key, Key): 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 +866,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 +888,12 @@ 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)) + # 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: