Make peekable() a decorator also

more_itertools/more.py

@@ -126,62 +126,10 @@ def first(iterable, default=_marker):
         return default
 
 
-class peekable(object):
-    """Wrap an iterator to allow lookahead and prepending elements.
-
-    Call :meth:`peek` on the result to get the value that will be returned
-    by :func:`next`. This won't advance the iterator:
-
-        >>> p = peekable(['a', 'b'])
-        >>> p.peek()
-        'a'
-        >>> next(p)
-        'a'
-
-    Pass :meth:`peek` a default value to return that instead of raising
-    ``StopIteration`` when the iterator is exhausted.
-
-        >>> p = peekable([])
-        >>> p.peek('hi')
-        'hi'
-
-    peekables also offer a :meth:`prepend` method, which "inserts" items
-    at the head of the iterable:
-
-        >>> p = peekable([1, 2, 3])
-        >>> p.prepend(10, 11, 12)
-        >>> next(p)
-        10
-        >>> p.peek()
-        11
-        >>> list(p)
-        [11, 12, 1, 2, 3]
-
-    peekables can be indexed. Index 0 is the item that will be returned by
-    :func:`next`, index 1 is the item after that, and so on:
-    The values up to the given index will be cached.
-
-        >>> p = peekable(['a', 'b', 'c', 'd'])
-        >>> p[0]
-        'a'
-        >>> p[1]
-        'b'
-        >>> next(p)
-        'a'
-
-    Negative indexes are supported, but be aware that they will cache the
-    remaining items in the source iterator, which may require significant
-    storage.
-
-    To check whether a peekable is exhausted, check its truth value:
-
-        >>> p = peekable(['a', 'b'])
-        >>> if p:  # peekable has items
-        ...     list(p)
-        ['a', 'b']
-        >>> if not p:  # peekable is exhaused
-        ...     list(p)
-        []
+class _Peekable(object):
+    """Wrapper class for iterables that allows for lookahead, indexing, and
+    other features . See :func:`peekable`, which allows this to be used as a
+    decorator for details.
 
     """
     def __init__(self, iterable):
@@ -198,15 +146,12 @@ def __bool__(self):
             return False
         return True
 
-    def __nonzero__(self):
-        # For Python 2 compatibility
-        return self.__bool__()
+    __nonzero__ = __bool__  # For Python 2 compatibility
 
     def peek(self, default=_marker):
-        """Return the item that will be next returned from ``next()``.
-
-        Return ``default`` if there are no items left. If ``default`` is not
-        provided, raise ``StopIteration``.
+        """Return the item that will be next returned from :meth:`next`,
+        or *default* if there are no items left. If *default* is not provided,
+        raise :exc:`StopIteration`.
 
         """
         if not self._cache:
@@ -219,33 +164,15 @@ def peek(self, default=_marker):
         return self._cache[0]
 
     def prepend(self, *items):
-        """Stack up items to be the next ones returned from ``next()`` or
-        ``self.peek()``. The items will be returned in
-        first in, first out order::
+        """"Insert" items at the head of the iterable. The items will be
+        inserted in first in, first out order:
 
             >>> p = peekable([1, 2, 3])
-            >>> p.prepend(10, 11, 12)
-            >>> next(p)
-            10
+            >>> p.prepend('a', 'b', 'c')
             >>> list(p)
-            [11, 12, 1, 2, 3]
-
-        It is possible, by prepending items, to "resurrect" a peekable that
-        previously raised ``StopIteration``.
-
-            >>> p = peekable([])
-            >>> next(p)
-            Traceback (most recent call last):
-              ...
-            StopIteration
-            >>> p.prepend(1)
-            >>> next(p)
-            1
-            >>> next(p)
-            Traceback (most recent call last):
-              ...
-            StopIteration
+            ['a', 'b', 'c', 1, 2, 3]
 
+        Items can be added to peekables that had previously been exhausted.
         """
         self._cache.extendleft(reversed(items))
 
@@ -296,6 +223,76 @@ def __getitem__(self, index):
         return self._cache[index]
 
 
+def peekable(iterable):
+    """Wrap an iterator to enable lookahead and sequence-like operations.
+
+    Call :meth:`peek` on a peekable-wrapped iterator to look ahead at the item
+    that will be returned by :func:`next`:
+
+        >>> p = peekable(['a', 'b'])
+        >>> p.peek()
+        'a'
+        >>> next(p)
+        'a'
+
+    Pass :meth:`peek` a default value to return that instead of raising
+    ``StopIteration`` when the iterator is exhausted:
+
+        >>> p = peekable([])
+        >>> p.peek('default')
+        'default'
+
+    To check whether a peekable is exhausted, check its truth value:
+
+        >>> p = peekable(['a', 'b'])
+        >>> if p:  # peekable has items
+        ...     list(p)
+        ['a', 'b']
+        >>> if not p:  # peekable is exhaused
+        ...     list(p)
+        []
+
+    Use :meth:`prepend` to insert items at the head of the peekable:
+
+        >>> p = peekable([0, 1, 2, 3])
+        >>> next(p)
+        0
+        >>> p.prepend('a', 'b', 'c')
+        >>> list(p)
+        ['a', 'b', 'c', 1, 2, 3]
+
+
+    peekables can be indexed. Index 0 is the item that will be returned by
+    :func:`next`, index 1 is the item after that, and so on.
+
+        >>> p = peekable(['a', 'b', 'c', 'd'])
+        >>> p[0]
+        'a'
+        >>> p[1]
+        'b'
+        >>> next(p), next(p)
+        ('a', 'b')
+
+    Indexing a peekable should behave like indexing a list, meaning both
+    negative indexes and slices are supported. Indexing will cache only the
+    necessary items, but be aware that this may require significant storage.
+
+    """
+    try:
+        iter(iterable)
+    except TypeError:
+        if not callable(iterable):
+            raise
+    else:
+        return _Peekable(iterable)
+
+    @wraps(iterable)
+    def peekable_wrapper(*args, **kwargs):
+        return _Peekable(iterable(*args, **kwargs))
+
+    return peekable_wrapper
+
+
 def _collate(*iterables, **kwargs):
     """Helper for ``collate()``, called when the user is using the ``reverse``
     or ``key`` keyword arguments on Python versions below 3.5.

more_itertools/tests/test_more.py

@@ -9,7 +9,7 @@
 from operator import add, itemgetter
 from unittest import TestCase
 
-from six.moves import filter, range, zip
+from six.moves import filter, map, range, zip
 
 import more_itertools as mi
 
@@ -102,10 +102,6 @@ def test_default(self):
 
 
 class PeekableTests(TestCase):
-    """Tests for ``peekable()`` behavor not incidentally covered by testing
-    ``collate()``
-
-    """
     def test_peek_default(self):
         """Make sure passing a default into ``peek()`` works."""
         p = mi.peekable([])
@@ -305,7 +301,6 @@ def test_prepend_slicing(self):
         p.prepend(30, 40, 50)
         pseq = [30, 40, 50] + seq  # pseq for prepended_seq
 
-        # adapt the specific tests from test_slicing
         self.assertEqual(p[0], 30)
         self.assertEqual(p[1:8], pseq[1:8])
         self.assertEqual(p[1:], pseq[1:])
@@ -362,6 +357,56 @@ def test_prepend_reversed(self):
         expected = [12, 11, 10, 0, 1, 2]
         self.assertEqual(actual, expected)
 
+    def test_decorate_function(self):
+        @mi.peekable
+        def generator_function(n):
+            """docstring"""
+            return map(str, range(n))
+
+        p = generator_function(5)
+        self.assertEqual(p.peek(), '0')
+        self.assertEqual(next(p), '0')
+        self.assertEqual(list(p), ['1', '2', '3', '4'])
+        self.assertEqual(generator_function.__doc__, 'docstring')
+
+    def test_decorate_method(self):
+        class Methodical(object):
+            def __init__(self, n):
+                self._n = n
+
+            @mi.peekable
+            def method(self, x):
+                """docstring"""
+                return map(str, range(self._n + x))
+
+        p = Methodical(1).method(4)
+        self.assertEqual(p.peek(), '0')
+        self.assertEqual(next(p), '0')
+        self.assertEqual(list(p), ['1', '2', '3', '4'])
+        self.assertEqual(Methodical(1).method.__doc__, 'docstring')
+
+    def test_decorate_class(self):
+        @mi.peekable
+        class Iterable(object):
+            """docstring"""
+
+            def __init__(self, n):
+                self._it = map(str, range(n))
+
+            def __iter__(self):
+                return self
+
+            def __next__(self):
+                return next(self._it)
+
+            next = __next__
+
+        p = Iterable(5)
+        self.assertEqual(p.peek(), '0')
+        self.assertEqual(next(p), '0')
+        self.assertEqual(list(p), ['1', '2', '3', '4'])
+        self.assertEqual(Iterable.__doc__, 'docstring')
+
 
 class ConsumerTests(TestCase):
     """Tests for ``consumer()``"""
@@ -1620,6 +1665,8 @@ def __iter__(self):
     def test_decorate_iterable_class(self):
         @mi.seekable
         class Iterable(object):
+            """docstring"""
+
             def __init__(self, n):
                 self._it = (i for i in range(n))
 
@@ -1636,6 +1683,7 @@ def __next__(self):
         self.assertEqual(list(s), [0, 1, 2, 3, 4])
         s.seek(1)
         self.assertEqual(list(s), [1, 2, 3, 4])
+        self.assertEqual(Iterable.__doc__, 'docstring')
 
     def test_decorate_callable_class(self):
         @mi.seekable