From d559709d4b97ca44c17a5bc707632a38f0c0122b Mon Sep 17 00:00:00 2001 From: hxperl Date: Fri, 11 Sep 2026 14:13:41 +0900 Subject: [PATCH] Render each run of a split array-of-tables where it was written MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A document may interrupt an array of tables with an unrelated table and then continue it: [[fruit]] name = "apple" [settings] color = true [[fruit]] name = "banana" The parser collects only contiguous `[[fruit]]` headers into one AoT, so the second run arrives at `Container.append` as a separate AoT under a key that already exists. That branch merged its elements into the first AoT and returned, discarding where the run appeared, and rendering then emitted the whole array at the first run's position — moving `[settings]` below `[[fruit]] name = "banana"` on a plain parse/dump round-trip. Keep the elements in the single AoT, so `doc["fruit"]` is still the whole array, and record the later run with an `_AoTContinuation` marker in the body. Rendering splits the AoT across its runs at the recorded positions. Elements appended after parsing render with the last run, and a run whose elements have all been deleted renders nothing. test_parse_aot_without_ending_newline came from #422 (a missing final newline corrupting the dump) and pinned the reordered output of the day as its expectation; it now asserts the document is preserved, which is what that fix was after. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01DyaZ48KuBpkPpfPui9xCUK --- tests/test_toml_document.py | 107 +++++++++++++++++++++++++++++++----- tomlkit/container.py | 101 +++++++++++++++++++++++++++++++++- 2 files changed, 191 insertions(+), 17 deletions(-) diff --git a/tests/test_toml_document.py b/tests/test_toml_document.py index 7ba61f5e..fe22110a 100644 --- a/tests/test_toml_document.py +++ b/tests/test_toml_document.py @@ -1583,20 +1583,7 @@ def test_parse_aot_without_ending_newline() -> None: [[products]] name = "Nail"''' doc = parse(content) - assert ( - doc.as_string() - == """\ -[[products]] -name = "Hammer" - -[[products]] -name = "Nail" -[foo] - -[bar] - -""" - ) + assert doc.as_string() == content assert doc == { "products": [ {"name": "Hammer"}, @@ -1666,3 +1653,95 @@ def test_scalar_is_not_captured_by_table_rendered_from_dotted_key() -> None: doc["z"] = 2 assert doc.as_string() == "a.b = 1\nz = 2\n" + + +def test_split_array_of_tables_keeps_document_order() -> None: + content = """\ +[[fruit]] +name = "apple" + +[settings] +color = true + +[[fruit]] +name = "banana" + +[[fruit]] +name = "cherry" +""" + doc = parse(content) + + assert doc.as_string() == content + assert [table["name"] for table in doc["fruit"]] == ["apple", "banana", "cherry"] + assert doc.unwrap() == { + "fruit": [{"name": "apple"}, {"name": "banana"}, {"name": "cherry"}], + "settings": {"color": True}, + } + + +def test_split_array_of_tables_appends_to_the_last_run() -> None: + content = """\ +[[fruit]] +name = "apple" + +[settings] +color = true + +[[fruit]] +name = "banana" +""" + doc = parse(content) + table = tomlkit.table() + table["name"] = "cherry" + doc["fruit"].append(table) + + assert len(doc["fruit"]) == 3 + assert doc.as_string().index("cherry") > doc.as_string().index("color") + + +def test_split_array_of_tables_after_deleting_the_start_of_a_run() -> None: + content = """\ +[[fruit]] +name = "apple" + +[settings] +color = true + +[[fruit]] +name = "banana" + +[[fruit]] +name = "cherry" +""" + doc = parse(content) + del doc["fruit"][1] + + assert ( + doc.as_string() + == """\ +[[fruit]] +name = "apple" + +[settings] +color = true + +[[fruit]] +name = "cherry" +""" + ) + + +def test_split_array_of_tables_survives_a_copy() -> None: + content = """\ +[[fruit]] +name = "apple" + +[settings] +color = true + +[[fruit]] +name = "banana" +""" + doc = parse(content) + + assert copy.deepcopy(doc).as_string() == content diff --git a/tomlkit/container.py b/tomlkit/container.py index 8ff30d98..8aa258ee 100644 --- a/tomlkit/container.py +++ b/tomlkit/container.py @@ -32,6 +32,27 @@ _NOT_SET = object() +class _AoTContinuation(Null): + """Body entry standing in for a later run of an array-of-tables header. + + ``[[a]] ... [b] ... [[a]]`` is one array whose headers are written in two + runs. The elements live in a single :class:`AoT` so that ``doc["a"]`` is + that whole array; this marker records where a later run was written so it + renders back in place instead of being pulled up to the first run. + """ + + def __init__(self, key: Key, aot: AoT, tables: list[Table]) -> None: + super().__init__() + self.key = key + self.aot = aot + self.tables = tables + + def _getstate( # type: ignore[override] + self, protocol: int = 3 + ) -> tuple[Key, AoT, list[Table]]: + return self.key, self.aot, self.tables + + class Container(_CustomDict): # type: ignore[type-arg] """ A container for items within a TOMLDocument. @@ -53,6 +74,9 @@ def __init__(self, parsed: bool = False) -> None: # out-of-order tables doesn't have to scan every key in the map; # stale entries are filtered by the per-key isinstance check self._out_of_order_keys: set[Key] = set() + # whether any array of tables in this body is written as several runs; + # false for almost every document, and rendering skips a pass when so + self._has_aot_continuation = False @property def body(self) -> list[tuple[Key | None, Item]]: @@ -366,9 +390,20 @@ def append( # Tried to define an AoT after a table with the same name. raise KeyAlreadyPresent(key) + start = len(current.body) for table in item.body: current.append(table) + if self._parsed and len(current.body) > start: + self._has_aot_continuation = True + # A second run of ``[[key]]`` headers, separated from the + # first by an unrelated table. The elements all belong to + # the one array, but the run has to render back where it + # was written, so remember where it starts. + self._body.append( + (None, _AoTContinuation(key, current, current.body[start:])) + ) + return self else: raise KeyAlreadyPresent(key) @@ -633,7 +668,13 @@ def last_item(self) -> Item | None: def as_string(self) -> str: """Render as TOML string.""" s = "" + ranges = self._aot_render_ranges() if self._has_aot_continuation else {} for k, v in self._body: + if isinstance(v, _AoTContinuation): + start, end = ranges[id(v)] + if start != end: + s += self._render_aot(v.key, v.aot, body=v.aot.body[start:end]) + continue if k is not None: if isinstance(v, Table): if ( @@ -650,7 +691,9 @@ def as_string(self) -> str: and "\n" not in v.trivia.indent ): s += "\n" - s += self._render_aot(k, v) + aot_range = ranges.get(id(v)) + body = None if aot_range is None else v.body[slice(*aot_range)] + s += self._render_aot(k, v, body=body) else: s += self._render_simple_item(k, v) else: @@ -658,6 +701,49 @@ def as_string(self) -> str: return s + def _aot_render_ranges(self) -> dict[int, tuple[int, int]]: + """Map each split array-of-tables run to the elements it renders. + + An array of tables interrupted by an unrelated table is stored as a + single ``AoT`` plus one ``_AoTContinuation`` marker per later run. The + keys of the returned mapping are the ``id()`` of the ``AoT`` (its first + run) and of each marker. An empty mapping is the common case. + """ + markers: dict[int, list[_AoTContinuation]] = {} + + for _, v in self._body: + if isinstance(v, _AoTContinuation): + markers.setdefault(id(v.aot), []).append(v) + + ranges: dict[int, tuple[int, int]] = {} + for conts in markers.values(): + aot = conts[0].aot + position = {id(table): i for i, table in enumerate(aot.body)} + live: list[_AoTContinuation] = [] + bounds: list[int] = [0] + for cont in conts: + # Elements can have been deleted since parsing, so the run + # starts at the first of its elements that is still there. If + # none is, the run renders nothing. + start = next( + ( + position[id(table)] + for table in cont.tables + if id(table) in position + ), + None, + ) + if start is not None and start >= bounds[-1]: + live.append(cont) + bounds.append(start) + else: + ranges[id(cont)] = (0, 0) + bounds.append(len(aot.body)) + ranges[id(aot)] = (0, bounds[1]) + for n, cont in enumerate(live): + ranges[id(cont)] = (bounds[n + 1], bounds[n + 2]) + return ranges + def _render_table(self, key: Key, table: Table, prefix: str | None = None) -> str: cur = "" @@ -740,14 +826,20 @@ def _render_table(self, key: Key, table: Table, prefix: str | None = None) -> st return cur - def _render_aot(self, key: Key, aot: AoT, prefix: str | None = None) -> str: + def _render_aot( + self, + key: Key, + aot: AoT, + prefix: str | None = None, + body: list[Table] | None = None, + ) -> str: _key = key.as_string() if prefix is not None: _key = prefix + "." + _key cur = "" _key = decode(_key) - for table in aot.body: + for table in aot.body if body is None else body: cur += self._render_aot_table(table, prefix=_key) return cur @@ -1012,6 +1104,9 @@ def __setstate__(self, state: tuple[Any, ...]) -> None: self._out_of_order_keys = { k for k, v in self._map.items() if isinstance(v, tuple) } + self._has_aot_continuation = any( + isinstance(v, _AoTContinuation) for _, v in self._body + ) for key, item in self._body: if key is not None: