diff --git a/CHANGELOG.md b/CHANGELOG.md index df615099..28b5afa7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - Speed up repeated array appends by keeping existing indexes instead of rebuilding them. - Speed up creating string values and keys, e.g. when building a document from Python data and dumping it: strings are escaped in one pass instead of per character, and the escape and bare-key character sets are built once. +- Speed up deleting array items with `del`, `pop()` and `remove()`, and adding lines with `Array.add_line()`: only the affected indexes are updated instead of rebuilding them all, so emptying a long array with `pop()` or building one line by line is no longer quadratic. ### Fixed diff --git a/tests/test_items.py b/tests/test_items.py index 13853fb2..eebe2143 100644 --- a/tests/test_items.py +++ b/tests/test_items.py @@ -494,6 +494,23 @@ def test_array_slice_deletion_keeps_render_in_sync() -> None: assert a.as_string() == "[3]" +def test_array_index_deletion_keeps_later_items_addressable() -> None: + doc = parse("a = [\n # first\n 1,\n 2, # two\n # gap\n 3,\n 4,\n 5,\n]\n") + a = doc["a"] + + del a[1] + assert a.pop() == 5 + a.remove(3) + a[1] = 40 + a.insert(1, 30) + a.append(50) + + assert a == [1, 30, 40, 50] + assert doc.as_string() == ( + "a = [\n # first\n 1,\n # gap\n 30,\n 40,\n 50,\n]\n" + ) + + def test_array_multiline() -> None: t = item([1, 2, 3, 4, 5, 6, 7, 8]) t.multiline(True) @@ -666,6 +683,19 @@ def test_array_add_line() -> None: ) +def test_array_add_line_keeps_items_addressable() -> None: + doc = parse("a = [\n 1, # one\n]\n") + a = doc["a"] + a.add_line(2, 3, comment="two") + a.add_line(comment="note") + a.add_line(4) + a[2] = 30 + del a[0] + + assert a == [2, 30, 4] + assert doc.as_string() == "a = [\n 2, 30, # two\n # note\n 4,\n]\n" + + def test_array_add_line_single_line_comment_round_trips() -> None: # A single-line array whose last rendered element is a trailing comment # (from add_line(..., comment=...)) must keep the closing bracket off the diff --git a/tomlkit/items.py b/tomlkit/items.py index 4d56f564..e9dabee0 100644 --- a/tomlkit/items.py +++ b/tomlkit/items.py @@ -1655,13 +1655,16 @@ def add_line( new_values.append( Comment(Trivia(indent=indent, comment=f"# {comment}", trail="")) ) + length = len(self) list.extend(self, data_values) + groups = self._group_values(new_values) + pos = len(self._value) if len(self._value) > 0: last_item = self._value[-1] last_value_item = next( ( v - for v in self._value[::-1] + for v in reversed(self._value) if v.value is not None and not isinstance(v.value, Null) ), None, @@ -1669,12 +1672,17 @@ def add_line( if last_value_item is not None: last_value_item.comma = Whitespace(",") if last_item.is_whitespace(): - self._value[-1:-1] = self._group_values(new_values) - else: - self._value.extend(self._group_values(new_values)) + pos -= 1 + self._value[pos:pos] = groups + if len(self._index_map) == length: + # The new groups come after every existing value, so only the + # new values need an index. + for i, group in enumerate(groups, pos): + if group.value is not None and not isinstance(group.value, Null): + self._index_map[length] = i + length += 1 else: - self._value.extend(self._group_values(new_values)) - self._reindex() + self._reindex() def clear(self) -> None: """Clear the array.""" @@ -1836,7 +1844,16 @@ def __delitem__(self, key: int | slice) -> None: # type: ignore[override] # remove the comma of the last item v.comma = None - self._reindex() + index_map = self._index_map + if isinstance(key, slice) or len(index_map) != length: + self._reindex() + else: + # A single value group was removed: the positions after it now + # point one group earlier. Shifting only those keeps pop() and + # deletes near the end of a long array from rebuilding the map. + for i in range(indices_to_remove[0], length - 1): + index_map[i] = index_map[i + 1] - 1 + del index_map[length - 1] def _getstate(self, protocol: int = 3) -> tuple[list[Item], Trivia, bool]: return list(self._iter_items()), self._trivia, self._multiline