Fix sibling data loss when updating out-of-order tables

tests/test_toml_document.py

@@ -1603,6 +1603,28 @@ def test_delete_key_from_out_of_order_table() -> None:
     )
 
 
+@pytest.mark.parametrize("replacement", [{"z": 1}, [{"z": 1}], 1])
+@pytest.mark.parametrize("remove_subtable", [False, True])
+@pytest.mark.parametrize("sibling_first", [False, True])
+def test_update_out_of_order_table_preserves_siblings(
+    replacement: Any, remove_subtable: bool, sibling_first: bool
+) -> None:
+    doc = parse("[y.y.b]\nc = 1\n[y]\nd = 1\n[y.b]\ne = 2\n[y.y]\n")
+    if remove_subtable:
+        doc["y"]["y"].pop("b")
+
+    table = doc["y"]
+    updates = {"y": replacement, "b": {"e": -0.5}}
+    if sibling_first:
+        updates = dict(reversed(updates.items()))
+    table.update(updates)
+
+    expected = {"y": {"y": replacement, "d": 1, "b": {"e": -0.5}}}
+    assert table.unwrap() == expected["y"]
+    assert doc.unwrap() == expected
+    assert parse(doc.as_string()).unwrap() == expected
+
+
 def test_parse_aot_without_ending_newline() -> None:
     content = '''\
 [[products]]

tomlkit/container.py

@@ -1254,11 +1254,13 @@ def _is_table_or_aot(it: Any) -> bool:
         _key: Key = key if isinstance(key, Key) else SingleKey(key)
 
         if _key in self._tables_map:
-            # Overwrite the first table and remove others
+            # Overwrite the first occurrence and remove this key from later parts.
             map_indices = self._tables_map[_key]
             while len(map_indices) > 1:
                 table = self._tables[map_indices.pop()]
-                self._remove_table(table)
+                del table[key]
+                if not table:
+                    self._remove_table(table)
             old_value = self._tables[map_indices[0]][key]
             if _is_table_or_aot(old_value) and not _is_table_or_aot(value):
                 # Remove the entry from the map and set value again.