Merge pull request #244 from themiurgo/peek

Add peek function

AUTHORS.md

@@ -21,3 +21,5 @@ Tom Prince                                      [@tomprince](https://github.com/
 Bart van Merriënboer                            [@bartvm](https://github.com/bartvm)
 
 Nikolaos-Digenis Karagiannis                    [@digenis](https://github.com/digenis/)
+
+[Antonio Lima](https://twitter.com/themiurgo)   [@themiurgo](https://github.com/themiurgo/)

doc/source/api.rst

@@ -33,6 +33,7 @@ Itertoolz
    nth
    partition
    partition_all
+   peek
    pluck
    reduceby
    remove

toolz/itertoolz.py

@@ -12,7 +12,7 @@
            'first', 'second', 'nth', 'last', 'get', 'concat', 'concatv',
            'mapcat', 'cons', 'interpose', 'frequencies', 'reduceby', 'iterate',
            'sliding_window', 'partition', 'partition_all', 'count', 'pluck',
-           'join', 'tail', 'diff', 'topk')
+           'join', 'tail', 'diff', 'topk', 'peek')
 
 
 def remove(predicate, seq):
@@ -877,3 +877,22 @@ def topk(k, seq, key=None):
     if key and not callable(key):
         key = getter(key)
     return tuple(heapq.nlargest(k, seq, key=key))
+
+
+def peek(seq):
+    """ Retrieve the next element of a sequence
+
+    Returns the first element and an iterable equivalent to the original
+    sequence, still having the element retrieved.
+
+    >>> seq = [0, 1, 2, 3, 4]
+    >>> first, seq = peek(seq)
+    >>> first
+    0
+    >>> list(seq)
+    [0, 1, 2, 3, 4]
+
+    """
+    iterator = iter(seq)
+    item = next(iterator)
+    return item, itertools.chain([item], iterator)

toolz/tests/test_itertoolz.py

@@ -11,7 +11,7 @@
                              reduceby, iterate, accumulate,
                              sliding_window, count, partition,
                              partition_all, take_nth, pluck, join,
-                             diff, topk)
+                             diff, topk, peek)
 from toolz.compatibility import range, filter
 from operator import add, mul
 
@@ -458,3 +458,12 @@ def test_topk():
 
 def test_topk_is_stable():
     assert topk(4, [5, 9, 2, 1, 5, 3], key=lambda x: 1) == (5, 9, 2, 1)
+
+
+def test_peek():
+    alist = ["Alice", "Bob", "Carol"]
+    element, blist  = peek(alist)
+    element == alist[0]
+    assert list(blist) == alist
+
+    assert raises(StopIteration, lambda: peek([]))