Normalize array indices before updating formatted storage

tests/test_items.py

@@ -424,6 +424,118 @@ def test_array_slice_deletion_keeps_render_in_sync() -> None:
     assert a.as_string() == "[3]"
 
 
+class ArrayIndex:
+    def __init__(self, value: int) -> None:
+        self.value = value
+        self.calls = 0
+
+    def __index__(self) -> int:
+        self.calls += 1
+        if self.calls > 1:
+            raise AssertionError("index converted more than once")
+        return self.value
+
+
+@pytest.mark.parametrize("operation", ["set", "delete", "insert"])
+@pytest.mark.parametrize("position", [0, 1, -1, -3])
+def test_array_mutation_with_index_protocol(operation: str, position: int) -> None:
+    doc = parse("a = [\n    1, # first\n    2,\n    3,\n]\n")
+    a = doc["a"]
+    expected = [1, 2, 3]
+    index = ArrayIndex(position)
+
+    if operation == "set":
+        a[index] = 4
+        expected[position] = 4
+    elif operation == "delete":
+        del a[index]
+        del expected[position]
+    else:
+        a.insert(index, 4)
+        expected.insert(position, 4)
+
+    assert index.calls == 1
+    assert a == expected
+    assert parse(doc.as_string())["a"] == expected
+
+
+@pytest.mark.parametrize("position", [-10, 10])
+@pytest.mark.parametrize("operation", ["set", "delete", "insert"])
+def test_array_mutation_with_out_of_bounds_index(operation: str, position: int) -> None:
+    content = "a = [1, 2, 3]\n"
+    doc = parse(content)
+    a = doc["a"]
+    index = ArrayIndex(position)
+
+    if operation == "insert":
+        a.insert(index, 4)
+        expected = [1, 2, 3]
+        expected.insert(position, 4)
+        assert a == expected
+        assert parse(doc.as_string())["a"] == expected
+    else:
+        with pytest.raises(IndexError):
+            if operation == "set":
+                a[index] = 4
+            else:
+                del a[index]
+        assert a == [1, 2, 3]
+        assert doc.as_string() == content
+
+    assert index.calls == 1
+
+
+@pytest.mark.parametrize("operation", ["set", "delete", "insert"])
+def test_array_mutation_with_invalid_index_does_not_mutate(operation: str) -> None:
+    content = "a = [1, 2, 3]\n"
+    doc = parse(content)
+    a = doc["a"]
+
+    with pytest.raises(TypeError):
+        if operation == "set":
+            a[1.5] = 4
+        elif operation == "delete":
+            del a[1.5]
+        else:
+            a.insert(1.5, 4)
+
+    assert a == [1, 2, 3]
+    assert doc.as_string() == content
+
+
+@pytest.mark.parametrize("bounds", [(0, 3, 2), (-1, -4, -1), (-10, -4, -1)])
+def test_array_slice_deletion_converts_indices_once(
+    bounds: tuple[int, int, int],
+) -> None:
+    doc = parse("a = [1, 2, 3]\n")
+    a = doc["a"]
+    expected = [1, 2, 3]
+    indices = [ArrayIndex(value) for value in bounds]
+
+    del a[slice(*indices)]
+    del expected[slice(*bounds)]
+
+    assert [index.calls for index in indices] == [1, 1, 1]
+    assert a == expected
+    assert parse(doc.as_string())["a"] == expected
+
+
+def test_array_invalid_slice_does_not_mutate() -> None:
+    content = "a = [1, 2, 3]\n"
+    doc = parse(content)
+    a = doc["a"]
+    start = ArrayIndex(0)
+    step = ArrayIndex(0)
+
+    with pytest.raises(ValueError, match="slice step cannot be zero"):
+        del a[slice(start, None, step)]
+
+    assert step.calls == 1
+    assert start.calls == 0
+    assert a == [1, 2, 3]
+    assert doc.as_string() == content
+
+
 def test_array_multiline() -> None:
     t = item([1, 2, 3, 4, 5, 6, 7, 8])
     t.multiline(True)

tomlkit/items.py

@@ -4,6 +4,7 @@
 import copy
 import dataclasses
 import inspect
+import operator
 import re
 import string
 
@@ -19,6 +20,7 @@
 from enum import Enum
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import SupportsIndex
 from typing import TypeVar
 from typing import overload
 
@@ -1625,22 +1627,24 @@ def clear(self) -> None:
     def __len__(self) -> int:
         return list.__len__(self)
 
-    def item(self, index: int) -> Item:
+    def item(self, index: SupportsIndex) -> Item:
         return list.__getitem__(self, index)  # type: ignore[no-any-return]
 
-    def __getitem__(self, key: int | slice) -> Any:  # type: ignore[override]
+    def __getitem__(self, key: SupportsIndex | slice) -> Any:
         return list.__getitem__(self, key)
 
-    def __setitem__(self, key: int | slice, value: Any) -> None:  # type: ignore[override]
+    def __setitem__(self, key: SupportsIndex | slice, value: Any) -> None:
         if isinstance(key, slice):
             raise ValueError("slice assignment is not supported")
+        key = operator.index(key)
         it = item(value, _parent=self)
         list.__setitem__(self, key, it)
         if key < 0:
             key += len(self)
         self._value[self._index_map[key]].value = it
 
-    def insert(self, pos: int, value: Any) -> None:  # type: ignore[override]
+    def insert(self, pos: SupportsIndex, value: Any) -> None:
+        pos = operator.index(pos)
         it = item(value, _parent=self)
         length = len(self)
         if not isinstance(it, (Comment, Whitespace)):
@@ -1700,7 +1704,18 @@ def insert(self, pos: int, value: Any) -> None:  # type: ignore[override]
         self._value.insert(idx, new_item)
         self._reindex()
 
-    def __delitem__(self, key: int | slice) -> None:  # type: ignore[override]
+    def __delitem__(self, key: SupportsIndex | slice) -> None:
+        if isinstance(key, slice):
+            step = None if key.step is None else operator.index(key.step)
+            if step == 0:
+                raise ValueError("slice step cannot be zero")
+            key = slice(
+                None if key.start is None else operator.index(key.start),
+                None if key.stop is None else operator.index(key.stop),
+                step,
+            )
+        else:
+            key = operator.index(key)
         length = len(self)
         list.__delitem__(self, key)