Fix dotted-key child replacement capturing following siblings

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Change Log
 
+## [Unreleased]
+
+### Fixed
+
+- Replacing a dotted-key child with a table or AoT no longer captures following dotted siblings on round-trip. ([#556](https://github.com/python-poetry/tomlkit/issues/556))
+
 ## [0.15.1] - 2026-07-17
 
 ### Changed

tests/test_toml_document.py

@@ -1210,6 +1210,122 @@ def test_replace_value_with_table_keeps_following_dotted_sibling() -> None:
     assert parse(doc.as_string()) == {"c": {"d": 2}, "x": {}}
 
 
+def test_replace_dotted_key_child_with_table() -> None:
+    # https://github.com/python-poetry/tomlkit/issues/556
+    content = """a.b = 1
+a.c = 2
+a.d = 3
+"""
+    doc = parse(content)
+    doc["a"]["b"] = {"x": 9}
+    assert (
+        doc.as_string()
+        == """a.c = 2
+a.d = 3
+
+[a.b]
+x = 9
+"""
+    )
+    assert parse(doc.as_string()) == {"a": {"b": {"x": 9}, "c": 2, "d": 3}}
+
+    a = doc["a"]
+    a["e"] = 4
+    assert doc["a"]["e"] == 4
+    assert parse(doc.as_string()) == {"a": {"b": {"x": 9}, "c": 2, "d": 3, "e": 4}}
+
+
+def test_replace_dotted_key_child_with_empty_table() -> None:
+    # https://github.com/python-poetry/tomlkit/issues/556
+    content = """a.b = 1
+a.c = 2
+a.d = 3
+"""
+    doc = parse(content)
+    doc["a"]["b"] = {}
+    assert (
+        doc.as_string()
+        == """a.c = 2
+a.d = 3
+
+[a.b]
+"""
+    )
+    assert parse(doc.as_string()) == {"a": {"b": {}, "c": 2, "d": 3}}
+
+
+def test_replace_dotted_key_child_with_aot() -> None:
+    # https://github.com/python-poetry/tomlkit/issues/556
+    content = """a.b = 1
+a.c = 2
+a.d = 3
+"""
+    doc = parse(content)
+    doc["a"]["b"] = [{"x": 9}]
+    assert (
+        doc.as_string()
+        == """a.c = 2
+a.d = 3
+
+[[a.b]]
+x = 9
+"""
+    )
+    assert parse(doc.as_string()) == {"a": {"b": [{"x": 9}], "c": 2, "d": 3}}
+
+
+def test_replace_dotted_key_child_with_interleaved_unrelated_group() -> None:
+    content = """a.b = 1
+z.q = 1
+a.c = 2
+"""
+    doc = parse(content)
+    doc["a"]["b"] = {"x": 9}
+    assert doc["z"] == {"q": 1}
+    doc["z"]["q"] = 5
+    expected_data = {"z": {"q": 5}, "a": {"c": 2, "b": {"x": 9}}}
+    assert doc.unwrap() == expected_data
+    assert parse(doc.as_string()).unwrap() == expected_data
+
+
+def test_replace_dotted_key_child_with_interleaved_plain_key() -> None:
+    content = """a.b = 1
+m = 0
+a.c = 2
+"""
+    doc = parse(content)
+    doc["a"]["b"] = {"x": 9}
+    assert doc["m"] == 0
+    doc["m"] = 7
+    assert doc["m"] == 7
+    expected_data = {"m": 7, "a": {"c": 2, "b": {"x": 9}}}
+    assert doc.unwrap() == expected_data
+    assert parse(doc.as_string()).unwrap() == expected_data
+
+
+@pytest.mark.parametrize(
+    "replacement, expected_b",
+    [
+        ({}, {}),
+        ([{"x": 9}], [{"x": 9}]),
+    ],
+)
+def test_replace_dotted_key_child_interleaved_empty_table_and_aot(
+    replacement: Any, expected_b: Any
+) -> None:
+    content = """a.b = 1
+z.q = 1
+a.c = 2
+"""
+    doc = parse(content)
+    doc["a"]["b"] = replacement
+    assert doc["z"] == {"q": 1}
+    doc["z"]["q"] = 5
+    expected_data = {"z": {"q": 5}, "a": {"c": 2, "b": expected_b}}
+    assert doc.unwrap() == expected_data
+    assert parse(doc.as_string()).unwrap() == expected_data
+
+
 def test_replace_with_comment() -> None:
     content = 'a = "1"'
     doc = parse(content)

tomlkit/container.py

@@ -513,6 +513,27 @@ def remove(self, key: Key | str) -> Container:
 
         return self
 
+    def _reorder_body_entry(self, p_idx: int, target_pos: int, group_key: Key) -> None:
+        entry = self._body.pop(p_idx)
+        self._body.insert(target_pos - 1, entry)
+
+        new_map: dict[Key, int | tuple[int, ...]] = {}
+        for k in self._map:
+            idxs = [i for i, (b_k, _) in enumerate(self._body) if b_k == k]
+            if idxs:
+                new_map[k] = idxs[0] if len(idxs) == 1 else tuple(idxs)
+
+        for b_k, _ in self._body:
+            if b_k is not None and b_k not in new_map:
+                idxs = [j for j, (b_k2, _) in enumerate(self._body) if b_k2 == b_k]
+                new_map[b_k] = idxs[0] if len(idxs) == 1 else tuple(idxs)
+
+        self._map = new_map
+        self._out_of_order_keys = {
+            k for k, v in self._map.items() if isinstance(v, tuple)
+        }
+        self._validation_cache.clear()
+
     def _insert_after(
         self, key: Key | str, other_key: Key | str, item: Any
     ) -> Container:
@@ -1101,10 +1122,10 @@ def __init__(self, container: Container, indices: tuple[int, ...]) -> None:
                         self._internal_container._raw_append(k, v)
                     else:
                         v = merged
-                    key_indices = self._tables_map.setdefault(k, [])  # type: ignore[arg-type]
-                    if table_idx not in key_indices:
-                        key_indices.append(table_idx)
                     if k is not None:
+                        key_indices = self._tables_map.setdefault(k, [])
+                        if table_idx not in key_indices:
+                            key_indices.append(table_idx)
                         dict.__setitem__(self, k.key, v)
 
         self._internal_container._validate_out_of_order_table()
@@ -1192,6 +1213,92 @@ def _is_table_or_aot(it: Any) -> bool:
                 self[key] = value
                 return
             self._tables[map_indices[0]][key] = value
+            if not _is_table_or_aot(old_value) and _is_table_or_aot(value):
+                part_item = self._tables[map_indices[0]].item(key)
+                if isinstance(part_item, AoT) or (
+                    isinstance(part_item, Table)
+                    and (
+                        not part_item.is_super_table()
+                        or self._container._renders_table_header(part_item)
+                    )
+                ):
+                    group_key = next(
+                        (
+                            k
+                            for k, v in self._container._body
+                            if v is self._tables[map_indices[0]]
+                        ),
+                        None,
+                    )
+                    if group_key is not None:
+                        p_idx = -1
+                        last_inline_idx = -1
+                        first_later_header_idx = -1
+                        for idx, (k, v) in enumerate(self._container._body):
+                            if k == group_key:
+                                if v is self._tables[map_indices[0]]:
+                                    p_idx = idx
+                                else:
+                                    if p_idx != -1 and first_later_header_idx == -1:
+                                        first_later_header_idx = idx
+                            if (
+                                k is not None
+                                and k.key == group_key.key
+                                and k.is_dotted()
+                                and not (
+                                    isinstance(v, Table)
+                                    and self._container._renders_table_header(v)
+                                )
+                            ):
+                                last_inline_idx = idx
+
+                        if (
+                            p_idx != -1
+                            and last_inline_idx != -1
+                            and p_idx < last_inline_idx
+                        ):
+                            target_pos = (
+                                first_later_header_idx
+                                if (
+                                    first_later_header_idx != -1
+                                    and first_later_header_idx > last_inline_idx
+                                )
+                                else last_inline_idx + 1
+                            )
+                            self._container._reorder_body_entry(
+                                p_idx, target_pos, group_key
+                            )
+
+                            self._tables.clear()
+                            self._tables_map.clear()
+                            self._internal_container = Container(True)
+                            dict.clear(self)
+
+                            indices = self._container._map[group_key]
+                            if not isinstance(indices, tuple):
+                                indices = (indices,)
+
+                            for i in indices:
+                                _, _item = self._container._body[i]
+                                if isinstance(_item, Table):
+                                    self._tables.append(_item)
+                                    table_idx = len(self._tables) - 1
+                                    for k, v in _item.value.body:
+                                        merged = self._merge_aot_fragment(k, v)
+                                        if merged is None:
+                                            self._internal_container._raw_append(k, v)
+                                        else:
+                                            v = merged
+                                        if k is not None:
+                                            key_indices = self._tables_map.setdefault(
+                                                k, []
+                                            )
+                                            if table_idx not in key_indices:
+                                                key_indices.append(table_idx)
+                                            dict.__setitem__(self, k.key, v)
+
+                            self._internal_container._validate_out_of_order_table()
+                            return
         elif self._tables:
             if not _is_table_or_aot(value):  # if the value is a plain value
                 for table in self._tables: