From 22b44c99f12d6579dab92e6fd81f8c05fba05190 Mon Sep 17 00:00:00 2001 From: David PAVLOVSCHII Date: Tue, 28 Jul 2026 22:12:19 +0300 Subject: [PATCH] Speed up repeated scalar insertion Find the scalar/table boundary with a binary search instead of scanning either side in full, and update only mapped keys in the suffix shifted by insertion. Preserve the previous handling of null entries and non-fixed whitespace by classifying them with the next meaningful body item. --- CHANGELOG.md | 6 +++ tests/test_toml_document.py | 41 +++++++++++++++++ tomlkit/container.py | 88 +++++++++++++++++++++++++------------ 3 files changed, 108 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3737b80..1c53e87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Change Log +## [Unreleased] + +### Changed + +- Speed up repeated scalar insertion by binary-searching the scalar/table boundary and updating only mapped keys in the shifted suffix. ([#540](https://github.com/python-poetry/tomlkit/issues/540)) + ## [0.15.1] - 2026-07-17 ### Changed diff --git a/tests/test_toml_document.py b/tests/test_toml_document.py index 7ba61f5..a3f7b45 100644 --- a/tests/test_toml_document.py +++ b/tests/test_toml_document.py @@ -1666,3 +1666,44 @@ 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_appending_scalars_without_child_tables_keeps_trailing_spacing() -> None: + # https://github.com/python-poetry/tomlkit/issues/540 + doc = parse("[t]\n\n[next]\n") + table = doc["t"] + + for i in range(100): + table[f"k{i}"] = i + + assert doc.as_string().startswith("[t]\nk0 = 0\nk1 = 1\n") + assert "k99 = 99\n\n[next]\n" in doc.as_string() + assert doc.unwrap()["t"]["k99"] == 99 + + +@pytest.mark.parametrize( + ("child_header", "separator"), [("[t.a]", "\n\n"), ("[[t.a]]", "\n")] +) +def test_appending_scalars_before_child_table( + child_header: str, separator: str +) -> None: + # A child table keeps scalar values in the parent table's body. Find that + # boundary without scanning either the scalar or table region in full. + doc = parse(f"[t]\n{child_header}\n") + table = doc["t"] + + for i in range(100): + table[f"k{i}"] = i + + assert doc.as_string().startswith("[t]\nk0 = 0\nk1 = 1\n") + assert f"k99 = 99{separator}{child_header}\n" in doc.as_string() + assert parse(doc.as_string()).unwrap() == doc.unwrap() + + +def test_appending_scalar_before_whitespace_and_child_table() -> None: + doc = parse("[t]\na = 1\n\n[t.child]\nb = 2\n") + table = doc["t"] + + table["c"] = 3 + + assert doc.as_string() == "[t]\na = 1\nc = 3\n\n[t.child]\nb = 2\n" diff --git a/tomlkit/container.py b/tomlkit/container.py index 8ff30d9..ad92710 100644 --- a/tomlkit/container.py +++ b/tomlkit/container.py @@ -151,30 +151,43 @@ def _handle_dotted_key(self, key: Key, value: Item) -> None: return def _get_last_index_before_table(self) -> int: - last_index = -1 - for i, (k, v) in enumerate(self._body): - if isinstance(v, Null): - continue # Null elements are inserted after deletion - - if isinstance(v, Whitespace) and not v.is_fixed(): - continue - - if isinstance(v, (Table, AoT)) and k is not None and not k.is_dotted(): - break - - if ( - isinstance(v, Table) - and k is not None - and k.is_dotted() - and self._renders_table_header(v) + # Scalar values precede tables in the body, so find their boundary + # without scanning either region in full. + low = 0 + high = len(self._body) + while low < high: + mid = (low + high) // 2 + if self._is_table_region(mid): + high = mid + else: + low = mid + 1 + return low + + def _is_table_region(self, index: int) -> bool: + # Nulls and non-fixed whitespace do not define the boundary. Give them + # the classification of the next meaningful item so whitespace between + # scalars stays with the scalar region, while whitespace immediately + # before a table stays with the table region. Trailing ignored items are + # always after the last scalar and therefore belong to the table side. + while index < len(self._body): + key, value = self._body[index] + if not ( + isinstance(value, Null) + or (isinstance(value, Whitespace) and not value.is_fixed()) ): - # A dotted-key super table renders inline (`a.b = 1`) only as - # long as none of its children render a `[table]` header; once - # one does, anything appended after it would land inside that - # table's scope. - break - last_index = i - return last_index + 1 + return ( + isinstance(value, (Table, AoT)) + and key is not None + and ( + not key.is_dotted() + or ( + isinstance(value, Table) + and self._renders_table_header(value) + ) + ) + ) + index += 1 + return True def _renders_table_header(self, table: Table) -> bool: for k, v in table.value.body: @@ -393,7 +406,12 @@ def append( or "\n" in after_item.trivia.indent ): after_item.trivia.indent = "\n" + after_item.trivia.indent - return self._insert_at(last_index, key, item) + return self._insert_at( + last_index, + key, + item, + update_table_indices_only=not is_table, + ) else: previous_item = self._body[-1][1] if isinstance(previous_item, Table) and previous_item.is_super_table(): @@ -560,7 +578,13 @@ def _insert_after( return self - def _insert_at(self, idx: int, key: Key | str, item: Any) -> Container: + def _insert_at( + self, + idx: int, + key: Key | str, + item: Any, + update_table_indices_only: bool = False, + ) -> Container: if idx > len(self._body) - 1: raise ValueError(f"Unable to insert at position {idx}") @@ -579,8 +603,18 @@ def _insert_at(self, idx: int, key: Key | str, item: Any) -> Container: ): previous_item.trivia.trail += "\n" - # Increment indices after the current index - for k, v in self._map.items(): + # Values inserted at the scalar/table boundary only shift table keys; + # every scalar key is before ``idx`` already. Avoid visiting all prior + # scalar keys on every insertion, which made bulk insertion quadratic. + keys = ( + dict.fromkeys(k for k, _ in self._body[idx:] if k is not None) + if update_table_indices_only + else self._map + ) + for k in keys: + v = self._map.get(k) + if v is None: + continue if isinstance(v, tuple): new_indices = [] for v_ in v: