Track new keys in out-of-order table proxies

tests/test_toml_document.py

@@ -754,6 +754,50 @@ def test_out_of_order_table_merges_three_aot_fragments() -> None:
     assert hooks["state"]["z"] == 3
 
 
+@pytest.mark.parametrize("value", [4, "text", [1, 2], {"nested": 3}])
+@pytest.mark.parametrize("operation", ["delete", "pop"])
+def test_out_of_order_proxy_can_remove_new_key(value: Any, operation: str) -> None:
+    doc = parse("[a.b]\nx=1\n[other]\ny=2\n[a]\nexisting=3\n")
+    proxy = doc["a"]
+    proxy["new"] = value
+    assert proxy["new"] == value
+
+    if operation == "delete":
+        del proxy["new"]
+    else:
+        assert proxy.pop("new") == value
+
+    assert "new" not in proxy
+    assert "new" not in doc["a"]
+    assert parse(doc.as_string()).unwrap() == {
+        "a": {"b": {"x": 1}, "existing": 3},
+        "other": {"y": 2},
+    }
+
+
+@pytest.mark.parametrize(
+    ("initial", "replacement"),
+    [
+        (4, 5),
+        ({"old": 4}, {"new": 5}),
+        ({"old": 4}, 5),
+        (4, {"new": 5}),
+        ([{"old": 4}], 5),
+        (4, [{"new": 5}]),
+    ],
+)
+def test_out_of_order_proxy_can_replace_new_key(initial: Any, replacement: Any) -> None:
+    doc = parse("[a.b]\nx=1\n[other]\ny=2\n[a]\nexisting=3\n")
+    proxy = doc["a"]
+    proxy["new"] = initial
+    proxy["new"] = replacement
+    assert proxy["new"] == replacement
+    assert parse(doc.as_string()).unwrap() == {
+        "a": {"b": {"x": 1}, "existing": 3, "new": replacement},
+        "other": {"y": 2},
+    }
+
+
 def test_out_of_order_tables_are_still_dicts() -> None:
     content = """
 [a.a]

tomlkit/container.py

@@ -1268,11 +1268,12 @@ def _is_table_or_aot(it: Any) -> bool:
                 return
             self._tables[map_indices[0]][key] = value
         elif self._tables:
+            table_index = 0
             if not _is_table_or_aot(value):  # if the value is a plain value
-                for table in self._tables:
+                for i, table in enumerate(self._tables):
                     # find the first table that allows plain values
                     if any(not _is_table_or_aot(v) for _, v in table.items()):
-                        table[key] = value
+                        table_index = i
                         break
                 else:
                     # No part holds a plain value yet, so the chosen part must
@@ -1281,14 +1282,12 @@ def _is_table_or_aot(it: Any) -> bool:
                     # turning a header-less super part concrete here would emit a
                     # second, duplicate header next to an existing concrete part
                     # and produce TOML that no longer parses.
-                    for table in self._tables:
+                    for i, table in enumerate(self._tables):
                         if not table.is_super_table():
-                            table[key] = value
+                            table_index = i
                             break
-                    else:
-                        self._tables[0][key] = value
-            else:
-                self._tables[0][key] = value
+            self._tables[table_index][key] = value
+            self._tables_map[_key] = [table_index]
         else:
             self._container[key] = value