diff --git a/tests/test_toml_document.py b/tests/test_toml_document.py index 7ba61f5..923eede 100644 --- a/tests/test_toml_document.py +++ b/tests/test_toml_document.py @@ -1666,3 +1666,37 @@ 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_key_insertion_scaling() -> None: + # https://github.com/python-poetry/tomlkit/issues/540 + # Inserting many keys into a table should not exhibit quadratic performance. + doc = parse("[packages]\n") + packages = doc["packages"] + for i in range(1000): + packages[f"key_{i}"] = i + + assert len(packages) == 1000 + assert packages["key_0"] == 0 + assert packages["key_999"] == 999 + reparsed = parse(doc.as_string()) + assert len(reparsed["packages"]) == 1000 + assert reparsed["packages"]["key_500"] == 500 + + +def test_key_insertion_before_existing_tables() -> None: + doc = parse("[tbl]\nx = 1\n") + for i in range(5): + doc[f"a_{i}"] = i + + lines = doc.as_string().splitlines() + assert lines[:5] == [f"a_{i} = {i}" for i in range(5)] + assert "[tbl]" in lines + assert parse(doc.as_string()) == { + "a_0": 0, + "a_1": 1, + "a_2": 2, + "a_3": 3, + "a_4": 4, + "tbl": {"x": 1}, + } diff --git a/tomlkit/container.py b/tomlkit/container.py index 8ff30d9..683521f 100644 --- a/tomlkit/container.py +++ b/tomlkit/container.py @@ -44,6 +44,8 @@ def __init__(self, parsed: bool = False) -> None: self._body: list[tuple[Key | None, Item]] = [] self._parsed = parsed self._table_keys: list[Key] = [] + self._has_tables: bool = False + self._first_table_idx: int | None = None # number of already-validated fragments and the temp container they # were merged into, per out-of-order key; lets parse-time validation # resume where the previous pass stopped instead of re-merging every @@ -150,31 +152,59 @@ def _handle_dotted_key(self, key: Key, value: Item) -> None: self.append(name, table) return + def _is_table_header(self, k: Key | None, v: Item) -> bool: + if k is None: + return False + + if isinstance(v, (Table, AoT)) and not k.is_dotted(): + return True + + if isinstance(v, Table) and k.is_dotted() and self._renders_table_header(v): + # 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. + return True + + return False + def _get_last_index_before_table(self) -> int: - last_index = -1 - for i, (k, v) in enumerate(self._body): + if not self._has_tables: + end = len(self._body) + else: + end = None + if self._first_table_idx is not None: + if self._first_table_idx < len(self._body): + k, v = self._body[self._first_table_idx] + if self._is_table_header(k, v): + end = self._first_table_idx + else: + self._first_table_idx = None + else: + self._first_table_idx = None + + if end is None: + for i, (k, v) in enumerate(self._body): + if self._is_table_header(k, v): + self._first_table_idx = i + end = i + break + + if end is None: + self._has_tables = False + end = len(self._body) + + for i in range(end - 1, -1, -1): + _, v = self._body[i] if isinstance(v, Null): - continue # Null elements are inserted after deletion + continue 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 + return i + 1 - if ( - isinstance(v, Table) - and k is not None - and k.is_dotted() - and self._renders_table_header(v) - ): - # 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 0 def _renders_table_header(self, table: Table) -> bool: for k, v in table.value.body: @@ -467,6 +497,10 @@ def _raw_append(self, key: Key | None, item: Item) -> None: self._map[key] = len(self._body) self._body.append((key, item)) + if isinstance(item, (Table, AoT)): + self._has_tables = True + if self._first_table_idx is None and self._is_table_header(key, item): + self._first_table_idx = len(self._body) - 1 if item.is_table() and key is not None: self._table_keys.append(key) @@ -474,6 +508,8 @@ def _raw_append(self, key: Key | None, item: Item) -> None: dict.__setitem__(self, key.key, item.value) def _remove_at(self, idx: int) -> None: + if self._first_table_idx is not None and idx == self._first_table_idx: + self._first_table_idx = None key = self._body[idx][0] assert key is not None index = self._map.get(key) @@ -503,6 +539,13 @@ def remove(self, key: Key | str) -> Container: raise NonExistentKey(key) self._validation_cache.clear() + if self._first_table_idx is not None: + if isinstance(idx, tuple): + if self._first_table_idx in idx: + self._first_table_idx = None + elif idx == self._first_table_idx: + self._first_table_idx = None + if isinstance(idx, tuple): for i in idx: self._body[i] = (None, Null()) @@ -554,6 +597,13 @@ def _insert_after( self._map[other_key] = idx + 1 self._body.insert(idx + 1, (other_key, item)) + if isinstance(item, (Table, AoT)): + self._has_tables = True + if self._is_table_header(other_key, item): + if self._first_table_idx is None or (idx + 1) < self._first_table_idx: + self._first_table_idx = idx + 1 + elif self._first_table_idx is not None and (idx + 1) <= self._first_table_idx: + self._first_table_idx += 1 if key is not None: dict.__setitem__(self, other_key.key, item.value) @@ -602,6 +652,13 @@ def _insert_at(self, idx: int, key: Key | str, item: Any) -> Container: else: self._map[key] = idx self._body.insert(idx, (key, item)) + if isinstance(item, (Table, AoT)): + self._has_tables = True + if self._is_table_header(key, item): + if self._first_table_idx is None or idx < self._first_table_idx: + self._first_table_idx = idx + elif self._first_table_idx is not None and idx <= self._first_table_idx: + self._first_table_idx += 1 dict.__setitem__(self, key.key, item.value) @@ -873,6 +930,15 @@ def _replace_at( ) -> None: value = _item(value) self._validation_cache.clear() + if self._first_table_idx is not None: + if isinstance(idx, tuple): + if self._first_table_idx in idx: + self._first_table_idx = None + elif idx == self._first_table_idx: + self._first_table_idx = None + + if isinstance(value, (Table, AoT)): + self._has_tables = True if isinstance(idx, tuple): for i in idx[1:]: @@ -1009,6 +1075,8 @@ def __setstate__(self, state: tuple[Any, ...]) -> None: self._body = state[1] self._parsed = state[2] self._table_keys = state[3] + self._has_tables = any(isinstance(v, (Table, AoT)) for _, v in self._body) + self._first_table_idx = None self._out_of_order_keys = { k for k, v in self._map.items() if isinstance(v, tuple) } @@ -1027,6 +1095,8 @@ def __copy__(self) -> Self: c._body += self.body c._map.update(self._map) + c._has_tables = self._has_tables + c._first_table_idx = self._first_table_idx c._out_of_order_keys |= self._out_of_order_keys return c