Raise IterationError on StopIteration

toolz/exceptions.py

@@ -0,0 +1,6 @@
+
+__all__ = ('IterationError',)
+
+
+class IterationError(RuntimeError):
+    pass

toolz/itertoolz.py

@@ -7,6 +7,7 @@
 from random import Random
 from collections.abc import Sequence
 from toolz.utils import no_default
+from toolz.exceptions import IterationError
 
 
 __all__ = ('remove', 'accumulate', 'groupby', 'merge_sorted', 'interleave',
@@ -373,7 +374,9 @@ def first(seq):
     >>> first('ABC')
     'A'
     """
-    return next(iter(seq))
+    for rv in seq:
+        return rv
+    raise IterationError("Received empty sequence")
 
 
 def second(seq):
@@ -382,9 +385,10 @@ def second(seq):
     >>> second('ABC')
     'B'
     """
-    seq = iter(seq)
-    next(seq)
-    return next(seq)
+    try:
+        return first(itertools.islice(seq, 1, None))
+    except IterationError as exc:
+        raise IterationError("Lenth of seq is < 2") from exc
 
 
 def nth(n, seq):
@@ -396,7 +400,10 @@ def nth(n, seq):
     if isinstance(seq, (tuple, list, Sequence)):
         return seq[n]
     else:
-        return next(itertools.islice(seq, n, None))
+        try:
+            return first(itertools.islice(seq, n, None))
+        except IterationError as exc:
+            raise IterationError("Length of seq is < %d" % n) from exc
 
 
 def last(seq):
@@ -531,8 +538,11 @@ def interpose(el, seq):
     [1, 'a', 2, 'a', 3]
     """
     inposed = concat(zip(itertools.repeat(el), seq))
-    next(inposed)
-    return inposed
+    try:
+        next(inposed)
+        return inposed
+    except StopIteration:
+        raise IterationError("Received empty sequence")
 
 
 def frequencies(seq):
@@ -722,13 +732,16 @@ def partition_all(n, seq):
     """
     args = [iter(seq)] * n
     it = zip_longest(*args, fillvalue=no_pad)
+
     try:
         prev = next(it)
     except StopIteration:
         return
+
     for item in it:
         yield prev
         prev = item
+
     if prev[-1] is no_pad:
         try:
             # If seq defines __len__, then
@@ -997,8 +1010,11 @@ def peek(seq):
     [0, 1, 2, 3, 4]
     """
     iterator = iter(seq)
-    item = next(iterator)
-    return item, itertools.chain((item,), iterator)
+    try:
+        item = next(iterator)
+        return item, itertools.chain((item,), iterator)
+    except StopIteration:
+        raise IterationError("Received empty sequence")
 
 
 def peekn(n, seq):

toolz/tests/test_itertoolz.py

@@ -14,8 +14,10 @@
                              sliding_window, count, partition,
                              partition_all, take_nth, pluck, join,
                              diff, topk, peek, peekn, random_sample)
-from operator import add, mul
 
+from toolz.exceptions import IterationError
+
+from operator import add, mul
 
 # is comparison will fail between this and no_default
 no_default2 = loads(dumps('__no__default__'))
@@ -127,7 +129,7 @@ def test_nth():
     assert nth(2, iter('ABCDE')) == 'C'
     assert nth(1, (3, 2, 1)) == 2
     assert nth(0, {'foo': 'bar'}) == 'foo'
-    assert raises(StopIteration, lambda: nth(10, {10: 'foo'}))
+    assert raises(IterationError, lambda: nth(10, {10: 'foo'}))
     assert nth(-2, 'ABCDE') == 'D'
     assert raises(ValueError, lambda: nth(-2, iter('ABCDE')))
 
@@ -136,12 +138,14 @@ def test_first():
     assert first('ABCDE') == 'A'
     assert first((3, 2, 1)) == 3
     assert isinstance(first({0: 'zero', 1: 'one'}), int)
+    assert raises(IterationError, lambda: first([]))
 
 
 def test_second():
     assert second('ABCDE') == 'B'
     assert second((3, 2, 1)) == 2
     assert isinstance(second({0: 'zero', 1: 'one'}), int)
+    assert raises(IterationError, lambda: second([1]))
 
 
 def test_last():
@@ -228,6 +232,7 @@ def test_interpose():
     assert "tXaXrXzXaXn" == "".join(interpose("X", "tarzan"))
     assert list(interpose(0, itertools.repeat(1, 4))) == [1, 0, 1, 0, 1, 0, 1]
     assert list(interpose('.', ['a', 'b', 'c'])) == ['a', '.', 'b', '.', 'c']
+    assert raises(IterationError, lambda: interpose('a', []))
 
 
 def test_frequencies():
@@ -510,8 +515,7 @@ def test_peek():
     element, blist = peek(alist)
     assert element == alist[0]
     assert list(blist) == alist
-
-    assert raises(StopIteration, lambda: peek([]))
+    assert raises(IterationError, lambda: peek([]))
 
 
 def test_peekn():