Fix out-of-order AoT element extensions

CHANGELOG.md

@@ -33,6 +33,7 @@
 - Fix `comment()` producing invalid TOML for a multiline string by prefixing every line with `#`, not just the first. ([#449](https://github.com/python-poetry/tomlkit/issues/449))
 - Fix the separator comma being swallowed by a trailing comment when appending a key to a multiline inline table, leaving the new key without a separator so the result no longer round-trips. ([#512](https://github.com/python-poetry/tomlkit/issues/512))
 - Fix a `KeyAlreadyPresent` error when parsing or accessing an out-of-order table whose array-of-tables elements are split across the table's parts. ([#505](https://github.com/python-poetry/tomlkit/issues/505))
+- Fix a `KeyAlreadyPresent` error when parsing, accessing, or unwrapping an out-of-order table whose later part extends the last element of an array of tables. ([#577](https://github.com/python-poetry/tomlkit/issues/577))
 - Out-of-order value-vs-table and dotted-key-vs-table redefinitions are now rejected at parse time instead of being silently accepted or raising only on access. The parser also detects when a non-dotted key is a prefix of an existing dotted key, matching the stdlib `tomllib` behaviour. ([#523](https://github.com/python-poetry/tomlkit/issues/523))
 - Reject tables inserted into inline tables instead of serializing invalid TOML. ([#531](https://github.com/python-poetry/tomlkit/issues/531))
 - Fix assigning an array of tables over a dotted key (e.g. `doc["a"] = aot(...)` where `a` came from `a.b = ...`): the new `[[a]]` header kept the dotted key's inline position and swallowed the following dotted sibling on round-trip. The array of tables now renders past the inline entries it would otherwise capture, mirroring the table fix for [#513](https://github.com/python-poetry/tomlkit/issues/513). ([#542](https://github.com/python-poetry/tomlkit/issues/542))

tests/test_toml_document.py

@@ -722,6 +722,53 @@ def test_out_of_order_table_merges_three_aot_fragments() -> None:
     assert hooks["state"]["z"] == 3
 
 
+def test_out_of_order_table_merges_aot_element_extension() -> None:
+    content = """\
+[[y.a]]
+name = "first"
+
+[b.d.x]
+
+[[y.a.c]]
+kind = "child"
+
+[[y.a]]
+name = "second"
+"""
+    doc = parse(content)
+    assert doc.as_string() == content
+    assert doc.unwrap() == {
+        "y": {
+            "a": [
+                {"name": "first", "c": [{"kind": "child"}]},
+                {"name": "second"},
+            ]
+        },
+        "b": {"d": {"x": {}}},
+    }
+
+    element = doc["y"]["a"][0]
+    element["name"] = "patched"
+    element["c"][0]["kind"] = "patched-child"
+    element["added"] = 3
+
+    expected = """\
+[[y.a]]
+name = "patched"
+added = 3
+
+[b.d.x]
+
+[[y.a.c]]
+kind = "patched-child"
+
+[[y.a]]
+name = "second"
+"""
+    assert doc.as_string() == expected
+    assert parse(expected).unwrap() == doc.unwrap()
+
+
 def test_out_of_order_tables_are_still_dicts() -> None:
     content = """
 [a.a]

tomlkit/container.py

@@ -304,7 +304,8 @@ def append(
                     if item.is_super_table():
                         # We need to merge both super tables
                         if (
-                            key.is_dotted()
+                            isinstance(current_idx, tuple)
+                            or key.is_dotted()
                             or (
                                 current_body_element[0] is not None
                                 and current_body_element[0].is_dotted()
@@ -1111,7 +1112,7 @@ def __init__(self, container: Container, indices: tuple[int, ...]) -> None:
 
     def _merge_aot_fragment(self, key: Key | None, item: Item) -> AoT | None:
         """
-        Merge an array-of-tables fragment from a later out-of-order table part.
+        Merge an array-of-tables fragment or extension from a later table part.
 
         An AoT whose elements are split across out-of-order parts of the same
         table arrives here once per part; ``_raw_append`` only knows how to
@@ -1120,10 +1121,15 @@ def _merge_aot_fragment(self, key: Key | None, item: Item) -> AoT | None:
         live element tables, without mutating either fragment (the parts keep
         rendering their own elements).
 
+        A later super table can also extend the last element of an AoT, as in
+        ``[[x.items]]`` followed by ``[[x.items.children]]``. In that case the
+        last element is presented through another out-of-order proxy spanning
+        the live AoT element and its later extension.
+
         Returns the merged ``AoT``, or ``None`` if this is not such a fragment.
         """
         internal = self._internal_container
-        if key is None or not isinstance(item, AoT) or key not in internal._map:
+        if key is None or key not in internal._map:
             return None
         idx = internal._map[key]
         if isinstance(idx, tuple):
@@ -1137,11 +1143,71 @@ def _merge_aot_fragment(self, key: Key | None, item: Item) -> AoT | None:
         if not isinstance(existing, AoT):
             return None
 
-        merged = AoT([*existing.body, *item.body], parsed=True)
+        if isinstance(item, AoT):
+            merged = AoT([*existing.body, *item.body], parsed=True)
+        elif isinstance(item, Table) and item.is_super_table() and existing.body:
+            # Container.append applies this same TOML rule for adjacent parts by
+            # mutating the last AoT element. These are out-of-order parts, so
+            # mutating either live source would make the document render the
+            # later header twice. Build a two-part proxy instead: reads merge
+            # both parts and writes still route to the original source table.
+            last = existing.body[-1]
+            fragments = Container(True)
+            fragment_key = SingleKey("__aot_element")
+            fragments._raw_append(fragment_key, last)
+            fragments._raw_append(fragment_key, item)
+            value = OutOfOrderTableProxy(fragments, (0, 1))
+            merged_last = Table(
+                value,  # type: ignore[arg-type]
+                copy.copy(last.trivia),
+                last.is_aot_element(),
+                last.is_super_table(),
+                last.name,
+                last.display_name,
+            )
+            merged = AoT([*existing.body[:-1], merged_last], parsed=True)
+        else:
+            return None
+
         internal._body[idx] = (internal._body[idx][0], merged)
         dict.__setitem__(internal, key.key, merged.value)
         return merged
 
+    @property
+    def body(self) -> list[tuple[Key | None, Item]]:
+        return self._internal_container.body
+
+    @property
+    def _map(self) -> dict[Key, int | tuple[int, ...]]:
+        return self._internal_container._map
+
+    def as_string(self) -> str:
+        return self._internal_container.as_string()
+
+    def copy(self) -> Container:
+        return self._internal_container.copy()
+
+    def append(
+        self, key: Key | str | None, item: Any, validate: bool = True
+    ) -> OutOfOrderTableProxy:
+        if key is None:
+            item = _item(item)
+            target = self._tables[0].value if self._tables else self._container
+            target.append(None, item, validate=validate)
+            self._internal_container.append(None, item, validate=validate)
+        else:
+            self[key] = item
+        return self
+
+    def remove(self, key: Key | str) -> OutOfOrderTableProxy:
+        del self[key]
+        return self
+
+    def _previous_item(
+        self, idx: int | None = None, ignore: tuple[type, ...] = (Null,)
+    ) -> Item | None:
+        return self._internal_container._previous_item(idx, ignore)
+
     def unwrap(self) -> dict[str, Any]:
         return self._internal_container.unwrap()