Honor the index protocol in positional operations

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')
 

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."

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()