Added peekn

doc/source/api.rst

@@ -34,6 +34,7 @@ Itertoolz
    partition
    partition_all
    peek
+   peekn
    pluck
    random_sample
    reduceby

toolz/curried/__init__.py

@@ -80,6 +80,7 @@
 partition = toolz.curry(toolz.partition)
 partition_all = toolz.curry(toolz.partition_all)
 partitionby = toolz.curry(toolz.partitionby)
+peekn = toolz.curry(toolz.peekn)
 pluck = toolz.curry(toolz.pluck)
 random_sample = toolz.curry(toolz.random_sample)
 reduce = toolz.curry(toolz.reduce)

toolz/itertoolz.py

@@ -14,7 +14,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', 'peek', 'random_sample')
+           'join', 'tail', 'diff', 'topk', 'peek', 'peekn', 'random_sample')
 
 
 def remove(predicate, seq):
@@ -942,7 +942,25 @@ def peek(seq):
     """
     iterator = iter(seq)
     item = next(iterator)
-    return item, itertools.chain([item], iterator)
+    return item, itertools.chain((item,), iterator)
+
+
+def peekn(n, seq):
+    """ Retrieve the next n elements of a sequence
+
+    Returns a tuple of the first n elements and an iterable equivalent
+    to the original, still having the elements retrieved.
+
+    >>> seq = [0, 1, 2, 3, 4]
+    >>> first_two, seq = peekn(2, seq)
+    >>> first_two
+    (0, 1)
+    >>> list(seq)
+    [0, 1, 2, 3, 4]
+    """
+    iterator = iter(seq)
+    peeked = tuple(take(n, iterator))
+    return peeked, itertools.chain(iter(peeked), iterator)
 
 
 def random_sample(prob, seq, random_state=None):

toolz/tests/test_itertoolz.py

@@ -13,7 +13,7 @@
                              reduceby, iterate, accumulate,
                              sliding_window, count, partition,
                              partition_all, take_nth, pluck, join,
-                             diff, topk, peek, random_sample)
+                             diff, topk, peek, peekn, random_sample)
 from toolz.compatibility import range, filter
 from operator import add, mul
 
@@ -496,12 +496,23 @@ def test_topk_is_stable():
 def test_peek():
     alist = ["Alice", "Bob", "Carol"]
     element, blist = peek(alist)
-    element == alist[0]
+    assert element == alist[0]
     assert list(blist) == alist
 
     assert raises(StopIteration, lambda: peek([]))
 
 
+def test_peekn():
+    alist = ("Alice", "Bob", "Carol")
+    elements, blist = peekn(2, alist)
+    assert elements == alist[:2]
+    assert tuple(blist) == alist
+
+    elements, blist = peekn(len(alist) * 4, alist)
+    assert elements == alist
+    assert tuple(blist) == alist
+
+
 def test_random_sample():
     alist = list(range(100))