diff --git a/tests/test_toml_document.py b/tests/test_toml_document.py index 7ba61f5..fe22110 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 8ff30d9..8aa258e 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: