From fc1bf69f877b07c6b72a8249e072dc4bd1ff1f5d Mon Sep 17 00:00:00 2001 From: repzakske <31244169+repzakske@users.noreply.github.com> Date: Thu, 15 Oct 2026 20:56:48 +0200 Subject: [PATCH] Honor the index protocol in positional operations --- src/sortedcontainers/sortedlist.py | 4 ++ src/sortedcontainers/sortedset.py | 4 +- tests/test_index_protocol.py | 98 ++++++++++++++++++++++++++++++ 3 files changed, 104 insertions(+), 2 deletions(-) create mode 100644 tests/test_index_protocol.py diff --git a/src/sortedcontainers/sortedlist.py b/src/sortedcontainers/sortedlist.py index aecf787..3ac0bc1 100644 --- a/src/sortedcontainers/sortedlist.py +++ b/src/sortedcontainers/sortedlist.py @@ -23,6 +23,7 @@ from itertools import chain, repeat, starmap from math import log from operator import add, eq, ge, gt, iadd, le, lt, ne +from operator import index as to_index from reprlib import recursive_repr from textwrap import dedent @@ -739,6 +740,7 @@ def __delitem__(self, index): pos, idx = _pos(index) _delete(pos, idx) else: + index = to_index(index) pos, idx = self._pos(index) self._delete(pos, idx) @@ -811,6 +813,7 @@ def __getitem__(self, index): indices = range(start, stop, step) return list(self._getitem(index) for index in indices) else: + index = to_index(index) if self._len: if index == 0: return _lists[0][0] @@ -1251,6 +1254,7 @@ def pop(self, index=-1): :raises IndexError: if index is out of range """ + index = to_index(index) if not self._len: raise IndexError('pop index out of range') diff --git a/src/sortedcontainers/sortedset.py b/src/sortedcontainers/sortedset.py index f2416bf..434a8ad 100644 --- a/src/sortedcontainers/sortedset.py +++ b/src/sortedcontainers/sortedset.py @@ -244,10 +244,10 @@ def __delitem__(self, index): if isinstance(index, slice): values = _list[index] _set.difference_update(values) + del _list[index] else: - value = _list[index] + value = _list.pop(index) _set.remove(value) - del _list[index] def __make_cmp(set_op, symbol, doc): "Make comparator method." diff --git a/tests/test_index_protocol.py b/tests/test_index_protocol.py new file mode 100644 index 0000000..209caa9 --- /dev/null +++ b/tests/test_index_protocol.py @@ -0,0 +1,98 @@ +from operator import delitem, getitem + +import pytest + +from sortedcontainers import SortedDict, SortedList, SortedSet + + +class Index: + def __init__(self, value): + self.value = value + self.calls = 0 + + def __index__(self): + self.calls += 1 + return self.value + + +class IntOnly: + def __int__(self): + return 0 + + +def pop_at(values, index): + return values.pop(index) + + +@pytest.mark.parametrize('key', [None, abs]) +@pytest.mark.parametrize('operation', [getitem, delitem, pop_at]) +@pytest.mark.parametrize('position', [0, 1, 15, 29, -1, -15, -30]) +def test_scalar_index_protocol(key, operation, position): + expected = list(range(30)) + values = SortedList(expected, key=key) + values._reset(4) + index = Index(position) + + assert operation(values, index) == operation(expected, position) + assert index.calls == 1 + assert list(values) == expected + values._check() + + +@pytest.mark.parametrize('key', [None, abs]) +@pytest.mark.parametrize('operation', [getitem, delitem, pop_at]) +@pytest.mark.parametrize('size', [0, 30]) +@pytest.mark.parametrize('index', [0.0, -1.0, 1.0, IntOnly()]) +def test_invalid_scalar_indices_do_not_modify_list(key, operation, size, index): + values = SortedList(range(size), key=key) + values._reset(4) + with pytest.raises(TypeError): + operation(values, index) + assert list(values) == list(range(size)) + values._check() + + +@pytest.mark.parametrize('key', [None, abs]) +@pytest.mark.parametrize('operation', [getitem, delitem, pop_at]) +@pytest.mark.parametrize('size', [0, 30]) +@pytest.mark.parametrize('position', [30, -31, 10**100]) +def test_index_protocol_out_of_range(key, operation, size, position): + values = SortedList(range(size), key=key) + index = Index(position) + with pytest.raises(IndexError): + operation(values, index) + assert index.calls == 1 + assert list(values) == list(range(size)) + values._check() + + +@pytest.mark.parametrize('operation', [getitem, delitem, pop_at]) +@pytest.mark.parametrize('index', [False, True]) +def test_boolean_indices(operation, index): + expected = [1, 2, 3] + values = SortedList(expected) + assert operation(values, index) == operation(expected, index) + assert list(values) == expected + values._check() + + +def test_sorted_set_deletion_converts_index_once(): + values = SortedSet(range(5)) + index = Index(2) + del values[index] + assert index.calls == 1 + assert list(values) == [0, 1, 3, 4] + assert 2 not in values + values._check() + + +def test_sorted_dict_positional_operations(): + values = SortedDict({1: 'a', 2: 'b', 3: 'c'}) + assert values.peekitem(Index(0)) == (1, 'a') + assert values.keys()[Index(1)] == 2 + assert values.items()[Index(1)] == (2, 'b') + assert values.values()[Index(1)] == 'b' + assert values.popitem(Index(-1)) == (3, 'c') + del values.keys()[Index(0)] + assert list(values.items()) == [(2, 'b')] + values._check()