Merge 9c113f2 into b2c747c

more_itertools/more.py

@@ -58,6 +58,7 @@
     'make_decorator',
     'map_except',
     'map_reduce',
+    'nth_or_last',
     'numeric_range',
     'one',
     'only',
@@ -180,6 +181,23 @@ def last(iterable, default=_marker):
         return default
 
 
+def nth_or_last(iterable, n, default=_marker):
+    """Return the nth or the last item of *iterable*,
+    or *default* if *iterable* is empty.
+
+        >>> nth_or_last([0, 1, 2, 3], 2)
+        2
+        >>> nth_or_last([0, 1], 2)
+        1
+        >>> nth_or_last([], 0, 'some default')
+        'some default'
+
+    If *default* is not provided and there are no items in the iterable,
+    raise ``ValueError``.
+    """
+    return last(islice(iterable, n + 1), default=default)
+
+
 class peekable:
     """Wrap an iterator to allow lookahead and prepending elements.
 

more_itertools/more.pyi

@@ -44,6 +44,11 @@ def last(iterable: Iterable[_T]) -> _T: ...
 @overload
 def last(iterable: Iterable[_T], default: _U) -> Union[_T, _U]: ...
 
+@overload
+def nth_or_last(iterable: Iterable[_T], n: int) -> _T: ...
+@overload
+def nth_or_last(iterable: Iterable[_T], n: int, default: _U) -> Union[_T, _U]: ...
+
 class peekable(Generic[_T], Iterator[_T]):
     def __init__(self, iterable: Iterable[_T]) -> None: ...
     def __iter__(self) -> peekable[_T]: ...

tests/test_more.py

@@ -200,6 +200,21 @@ def test_customrange(self):
         self.assertEqual(mi.last(IterOnlyRange(5)), 4)
 
 
+class NthOrLastTests(TestCase):
+    """Tests for ``nth_or_last()``"""
+
+    def test_basic(self):
+        self.assertEqual(mi.nth_or_last(range(3), 1), 1)
+        self.assertEqual(mi.nth_or_last(range(3), 3), 2)
+
+    def test_default_value(self):
+        default = 42
+        self.assertEqual(mi.nth_or_last(range(0), 3, default), default)
+
+    def test_empty_iterable_no_default(self):
+        self.assertRaises(ValueError, lambda: mi.nth_or_last(range(0), 0))
+
+
 class PeekableTests(TestCase):
     """Tests for ``peekable()`` behavor not incidentally covered by testing
     ``collate()``