add support for negative index to element_at()

reflect .net6 ElementAt's index overload

Author: cleoold

doc/api/types_linq.enumerable.rst

@@ -639,6 +639,8 @@ Returns
 Returns the element at specified index in the sequence. `IndexOutOfRangeError` is raised if
 no such element exists.
 
+If the index is negative, it means counting from the end.
+
 This method always uses a generic list element-finding method (O(n)) regardless the
 implementation of the wrapped iterable.
 
@@ -651,6 +653,9 @@ Example
         >>> Enumerable(gen()).element_at(1)
         10
 
+        >>> Enumerable(gen()).element_at(-1)
+        100
+
 ----
 
 instancemethod ``element_at[TDefault](index, __default)``
@@ -666,6 +671,8 @@ Returns
 Returns the element at specified index in the sequence. Default value is returned if no
 such element exists.
 
+If the index is negative, it means counting from the end.
+
 This method always uses a generic list element-finding method (O(n)) regardless the
 implementation of the wrapped iterable.
 

tests/test_usage.py

@@ -327,17 +327,27 @@ def test_overload1_0(self):
         gen = lambda: (i for i in range(7, 12))
         en = Enumerable(gen)
         assert en.element_at(0) == 7
+        assert en.element_at(-5) == 7
 
     def test_overload1_end(self):
         gen = lambda: (i for i in range(7, 12))
         en = Enumerable(gen)
         assert en.element_at(4) == 11
+        assert en.element_at(-1) == 11
 
     def test_overload1_out(self):
         gen = lambda: (i for i in range(7, 12))
         en = Enumerable(gen)
         with pytest.raises(IndexOutOfRangeError):
             en.element_at(5)
+        with pytest.raises(IndexOutOfRangeError):
+            en.element_at(-6)
+
+    def test_overload2_within(self):
+        gen = lambda: (i for i in range(7, 12))
+        en = Enumerable(gen)
+        assert en.element_at(2, 0) == 9
+        assert en.element_at(-3, 0) == 9
 
     def test_overload2_out(self):
         gen = lambda: (i for i in range(7, 12))

types_linq/enumerable.py

@@ -55,6 +55,10 @@ def _contains_impl(self, value: object, fallback: bool) -> bool:
     def __contains__(self, value: object) -> bool:
         return self._contains_impl(value, fallback=False)
 
+    @staticmethod
+    def _raise_not_enough_elements() -> NoReturn:
+        raise IndexOutOfRangeError('Not enough elements in the sequence')
+
     def _every(self, step: int) -> Enumerable[TSource_co]:
         return self.where2(lambda _, i: i % step == 0)
 
@@ -71,20 +75,27 @@ def _getitem_impl(self,
                     return iterable[index]
                 except IndexError as e:
                     raise IndexOutOfRangeError from e
-            iterator = iter(iterable)
-            try:
-                for _ in range(index):
-                    next(iterator)
-                return next(iterator)
-            except StopIteration:
-                raise IndexOutOfRangeError('Not enough elements in the sequence')
+            if index >= 0:
+                iterator = iter(iterable)
+                try:
+                    for _ in range(index):
+                        next(iterator)
+                    return next(iterator)
+                except StopIteration:
+                    self._raise_not_enough_elements()
+            else:
+                en = iterable if isinstance(iterable, Enumerable) else Enumerable(iterable)
+                last = en.take_last(-index).to_list()
+                if len(last) < -index:
+                    self._raise_not_enough_elements()
+                return last[0]
 
         else:  # isinstance(index, slice)
             if not fallback and isinstance(iterable, Sequence):
                 try:
                     res = iterable[index]
                 except IndexError as e:
-                    raise IndexOutOfRangeError(e)
+                    raise IndexOutOfRangeError from e
                 return res if isinstance(res, Enumerable) else Enumerable(res)
             # we do not enumerate all values if the begin and the end only involve
             # nonnegative indices since in which case the sliced part can be obtained

types_linq/enumerable.pyi

@@ -476,6 +476,8 @@ class Enumerable(Sequence[TSource_co], Generic[TSource_co]):
         Returns the element at specified index in the sequence. `IndexOutOfRangeError` is raised if
         no such element exists.
 
+        If the index is negative, it means counting from the end.
+
         This method always uses a generic list element-finding method (O(n)) regardless the
         implementation of the wrapped iterable.
 
@@ -487,6 +489,9 @@ class Enumerable(Sequence[TSource_co], Generic[TSource_co]):
 
                 >>> Enumerable(gen()).element_at(1)
                 10
+
+                >>> Enumerable(gen()).element_at(-1)
+                100
         '''
 
     @overload
@@ -495,6 +500,8 @@ class Enumerable(Sequence[TSource_co], Generic[TSource_co]):
         Returns the element at specified index in the sequence. Default value is returned if no
         such element exists.
 
+        If the index is negative, it means counting from the end.
+
         This method always uses a generic list element-finding method (O(n)) regardless the
         implementation of the wrapped iterable.