Speed up key lookups and replacements in tables

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

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

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,24 +853,20 @@ 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):
             # during parsing every fragment append re-probes membership;
             # 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: