Fix promoted dotted-key header capturing following siblings (#556)

CHANGELOG.md

@@ -34,6 +34,7 @@
 - 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))
 - 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 replacing a dotted-key value with a table (e.g. `doc["a"]["b"] = {...}` where `a.b`/`a.c`/`a.d` were parsed as separate root fragments) capturing the sibling dotted keys that follow it. Promoting `a.b` to a `[a.b]` header used to leave the trailing root dotted keys (whether the same prefix `a.c`/`a.d` or a different one such as `p.q`) rendered after the header, so they were re-parsed as its children. Those inline entries now render before the promoted header, preserving the round-trip. ([#556](https://github.com/python-poetry/tomlkit/issues/556))
 
 ## [0.15.0] - 2026-05-10
 

tests/test_toml_document.py

@@ -1131,6 +1131,71 @@ 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_keeps_following_same_name_siblings() -> 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}
+    # Promoting ``a.b`` to a ``[a.b]`` header must not capture the sibling
+    # dotted keys ``a.c``/``a.d`` that render afterwards; they are emitted
+    # before the header so the round-trip is preserved.
+    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_middle_dotted_key_keeps_following_same_name_siblings() -> None:
+    # https://github.com/python-poetry/tomlkit/issues/556
+    content = """a.b = 1
+a.c = 2
+a.d = 3
+"""
+    doc = parse(content)
+    doc["a"]["c"] = {"x": 9}
+    # Only the sibling that renders after the promoted ``[a.c]`` header needs
+    # to move ahead of it; ``a.b`` already precedes the header and stays put.
+    assert (
+        doc.as_string()
+        == """a.b = 1
+a.d = 3
+
+[a.c]
+x = 9
+"""
+    )
+    assert parse(doc.as_string()) == {"a": {"b": 1, "c": {"x": 9}, "d": 3}}
+
+
+def test_replace_dotted_key_keeps_following_other_dotted_key() -> None:
+    # https://github.com/python-poetry/tomlkit/issues/556
+    # A promoted ``[a.b]`` header must not capture a differently-named dotted
+    # key (``p.q``) that renders after it either.
+    content = """a.b = 1
+p.q = 5
+"""
+    doc = parse(content)
+    doc["a"]["b"] = {"x": 9}
+    assert (
+        doc.as_string()
+        == """p.q = 5
+
+[a.b]
+x = 9
+"""
+    )
+    assert parse(doc.as_string()) == {"p": {"q": 5}, "a": {"b": {"x": 9}}}
+
+
 def test_replace_with_comment() -> None:
     content = 'a = "1"'
     doc = parse(content)

tomlkit/container.py

@@ -32,6 +32,29 @@
 _NOT_SET = object()
 
 
+def _dotted_fragment_emits_header(item: Item) -> bool:
+    """Whether rendering a dotted table fragment emits a ``[table]`` header.
+
+    Out-of-order fragments of the same dotted key (e.g. the three ``a`` parts
+    parsed from ``a.b``/``a.c``/``a.d``) are stored as separate super-table
+    ``Table`` items. A fragment that only holds scalars renders as dotted keys
+    (``a.c = 2``); one that (transitively) holds a real table or array-of-tables
+    renders a header line (``[a.b]``). A header emitted before a sibling dotted
+    fragment would visually capture that sibling, so callers use this to keep
+    the header-emitting fragments last -- see :meth:`Container._render_ordering`.
+    """
+    if not isinstance(item, Table) or not item.is_super_table():
+        return True
+    for _, v in item.value.body:
+        if isinstance(v, AoT):
+            return True
+        if isinstance(v, Table) and (
+            not v.is_super_table() or _dotted_fragment_emits_header(v)
+        ):
+            return True
+    return False
+
+
 class Container(_CustomDict):  # type: ignore[type-arg]
     """
     A container for items within a TOMLDocument.
@@ -606,10 +629,75 @@ def last_item(self) -> Item | None:
             return self._body[-1][1]
         return None
 
+    def _render_ordering(
+        self, body: list[tuple[Key | None, Item]]
+    ) -> list[tuple[Key | None, Item]]:
+        """Order dotted table headers after the inline entries they would capture.
+
+        When a dotted key (``a.b``) is replaced by a table it is promoted to a
+        ``[a.b]`` header, but dotted-key siblings that happen to sit later in the
+        body -- whether the same top-level name (``a.c``/``a.d``) or a different
+        one (``p.q``) -- still render as bare dotted keys. A ``[a.b]`` header
+        emitted before them would capture them as its children on re-parse
+        (https://github.com/python-poetry/tomlkit/issues/556). Within the leading
+        run of root entries (before the first real ``[table]``/array-of-tables
+        header), stably move every header-emitting dotted fragment after the
+        scalar-rendering entries that follow it, preserving relative order and
+        leaving everything from the first real header onward untouched. The
+        reordering only kicks in when a header-emitting dotted fragment precedes
+        such an entry, so ordinary documents are left as-is.
+        """
+
+        def kind(k: Key | None, v: Item) -> str:
+            if isinstance(v, AoT) or (
+                isinstance(v, Table) and k is not None and not k.is_dotted()
+            ):
+                return "header"  # real ``[table]`` / array-of-tables section
+            if (
+                isinstance(v, Table)
+                and k is not None
+                and k.is_dotted()
+                and _dotted_fragment_emits_header(v)
+            ):
+                return "promoted"  # dotted key promoted to a ``[a.b]`` header
+            if k is not None and not isinstance(v, Whitespace):
+                return "inline"  # renders as a value / dotted key at this level
+            return "other"  # whitespace, comments, deletion placeholders
+
+        kinds = [kind(k, v) for k, v in body]
+
+        # Only the leading root run (up to the first real header) can capture
+        # trailing inline entries; genuine sections keep their own ordering.
+        run_end = len(body)
+        for i, kd in enumerate(kinds):
+            if kd == "header":
+                run_end = i
+                break
+
+        promoted = [i for i in range(run_end) if kinds[i] == "promoted"]
+        last_inline = max(
+            (i for i in range(run_end) if kinds[i] == "inline"), default=-1
+        )
+        # Nothing to do unless a promoted header precedes an inline entry.
+        if not promoted or promoted[0] > last_inline:
+            return body
+
+        # Stable partition of the run: everything that is not a promoted header,
+        # in order, followed by the promoted headers, in order.
+        kept = [i for i in range(run_end) if kinds[i] != "promoted"]
+        reordered = kept + promoted
+        # ``kept`` and ``promoted`` partition ``range(run_end)``, so ``reordered``
+        # is a permutation of it with exactly ``run_end`` entries.
+        order = list(range(len(body)))
+        for pos, src in enumerate(reordered):
+            order[pos] = src
+
+        return [body[i] for i in order]
+
     def as_string(self) -> str:
         """Render as TOML string."""
         s = ""
-        for k, v in self._body:
+        for k, v in self._render_ordering(self._body):
             if k is not None:
                 if isinstance(v, Table):
                     if (