[pypy-commit] pypy length-hint: length_hint basics and some initial iterator support

[pjenvey]

Author: Philip Jenvey <pjenvey at underboss.org>
Branch: length-hint
Changeset: r54190:856f01d18746
Date: 2012-04-04 18:31 -0700
http://bitbucket.org/pypy/pypy/changeset/856f01d18746/

Log:	length_hint basics and some initial iterator support

diff --git a/pypy/objspace/std/iterobject.py b/pypy/objspace/std/iterobject.py
--- a/pypy/objspace/std/iterobject.py
+++ b/pypy/objspace/std/iterobject.py
@@ -4,6 +4,29 @@
 from pypy.objspace.std.register_all import register_all
 
 
+def length_hint(space, w_obj, default):
+    """Return the length of an object, consulting its __length_hint__
+    method if necessary.
+    """
+    try:
+        return space.len_w(w_obj)
+    except OperationError, e:
+        if not (e.match(space, space.w_TypeError) or
+                e.match(space, space.w_AttributeError)):
+            raise
+
+    try:
+        w_hint = space.call_method(w_obj, '__length_hint__')
+    except OperationError, e:
+        if not (e.match(space, space.w_TypeError) or
+                e.match(space, space.w_AttributeError)):
+            raise
+        return default
+
+    hint = space.int_w(w_hint)
+    return default if hint < 0 else hint
+
+
 class W_AbstractIterObject(W_Object):
     __slots__ = ()
 
@@ -71,10 +94,6 @@
     w_seqiter.index += 1
     return w_item
 
-# XXX __length_hint__()
-##def len__SeqIter(space,  w_seqiter):
-##    return w_seqiter.getlength(space)
-
 
 def iter__FastTupleIter(space, w_seqiter):
     return w_seqiter
@@ -92,10 +111,6 @@
     w_seqiter.index = index + 1
     return w_item
 
-# XXX __length_hint__()
-##def len__FastTupleIter(space, w_seqiter):
-##    return w_seqiter.getlength(space)
-
 
 def iter__FastListIter(space, w_seqiter):
     return w_seqiter
@@ -115,10 +130,6 @@
     w_seqiter.index = index + 1
     return w_item
 
-# XXX __length_hint__()
-##def len__FastListIter(space, w_seqiter):
-##    return w_seqiter.getlength(space)
-
 
 def iter__ReverseSeqIter(space, w_seqiter):
     return w_seqiter
@@ -136,20 +147,4 @@
         raise OperationError(space.w_StopIteration, space.w_None)
     return w_item
 
-# XXX __length_hint__()
-##def len__ReverseSeqIter(space, w_seqiter):
-##    if w_seqiter.w_seq is None:
-##        return space.wrap(0)
-##    index = w_seqiter.index+1
-##    w_length = space.len(w_seqiter.w_seq)
-##    # if length of sequence is less than index :exhaust iterator
-##    if space.is_true(space.gt(space.wrap(w_seqiter.index), w_length)):
-##        w_len = space.wrap(0)
-##        w_seqiter.w_seq = None
-##    else:
-##        w_len =space.wrap(index)
-##    if space.is_true(space.lt(w_len,space.wrap(0))):
-##        w_len = space.wrap(0)
-##    return w_len
-
 register_all(vars())
diff --git a/pypy/objspace/std/itertype.py b/pypy/objspace/std/itertype.py
--- a/pypy/objspace/std/itertype.py
+++ b/pypy/objspace/std/itertype.py
@@ -23,6 +23,12 @@
     tup      = [w_self.w_seq, space.wrap(w_self.index)]
     return space.newtuple([new_inst, space.newtuple(tup)])
 
+
+def descr_seqiter__length_hint__(space, w_self):
+    from pypy.objspace.std.iterobject import W_AbstractSeqIterObject
+    assert isinstance(w_self, W_AbstractSeqIterObject)
+    return w_self.getlength(space)
+
 # ____________________________________________________________
 
 def descr_reverseseqiter__reduce__(w_self, space):
@@ -39,6 +45,24 @@
     tup      = [w_self.w_seq, space.wrap(w_self.index)]
     return space.newtuple([new_inst, space.newtuple(tup)])
 
+
+def descr_reverseseqiter__length_hint__(space, w_self):
+    from pypy.objspace.std.iterobject import W_ReverseSeqIterObject
+    assert isinstance(w_self, W_ReverseSeqIterObject)
+    if w_self.w_seq is None:
+        return space.wrap(0)
+    index = w_self.index + 1
+    w_length = space.len(w_self.w_seq)
+    # if length of sequence is less than index :exhaust iterator
+    if space.is_true(space.gt(space.wrap(w_self.index), w_length)):
+        w_len = space.wrap(0)
+        w_self.w_seq = None
+    else:
+        w_len = space.wrap(index)
+    if space.is_true(space.lt(w_len, space.wrap(0))):
+        w_len = space.wrap(0)
+    return w_len
+
 # ____________________________________________________________
 iter_typedef = StdTypeDef("sequenceiterator",
     __doc__ = '''iter(collection) -> iterator
@@ -49,11 +73,13 @@
 In the second form, the callable is called until it returns the sentinel.''',
 
     __reduce__ = gateway.interp2app(descr_seqiter__reduce__),
+    __length_hint__ = gateway.interp2app(descr_seqiter__length_hint__),
     )
 iter_typedef.acceptable_as_base_class = False
 
 reverse_iter_typedef = StdTypeDef("reversesequenceiterator",
 
     __reduce__ = gateway.interp2app(descr_reverseseqiter__reduce__),
+    __length_hint__ = gateway.interp2app(descr_reverseseqiter__length_hint__),
     )
 reverse_iter_typedef.acceptable_as_base_class = False
diff --git a/pypy/objspace/std/test/test_lengthhint.py b/pypy/objspace/std/test/test_lengthhint.py
new file mode 100644
--- /dev/null
+++ b/pypy/objspace/std/test/test_lengthhint.py
@@ -0,0 +1,53 @@
+from pypy.objspace.std.iterobject import length_hint
+
+
+class TestLengthHint:
+
+    SIZE = 4
+    ITEMS = range(SIZE)
+
+    def _test_length_hint(self, w_obj):
+        space = self.space
+        assert length_hint(space, w_obj, 8) == self.SIZE
+
+        w_iter = space.iter(w_obj)
+        assert space.int_w(
+            space.call_method(w_iter, '__length_hint__')) == self.SIZE
+        assert length_hint(space, w_iter, 8) == self.SIZE
+
+        space.next(w_iter)
+        assert length_hint(space, w_iter, 8) == self.SIZE - 1
+
+    def test_list(self):
+        self._test_length_hint(self.space.newlist(self.ITEMS))
+
+    def test_tuple(self):
+        self._test_length_hint(self.space.newtuple(self.ITEMS))
+
+    def test_reversed(self):
+        space = self.space
+        w_reversed = space.call_method(space.builtin, 'reversed',
+                                       space.newlist(self.ITEMS))
+        assert space.int_w(
+            space.call_method(w_reversed, '__length_hint__')) == self.SIZE
+        self._test_length_hint(w_reversed)
+
+    def test_default(self):
+        space = self.space
+        assert length_hint(space, space.w_False, 3) == 3
+
+    def test_exc(self):
+        from pypy.interpreter.error import OperationError
+        space = self.space
+        w_foo = space.appexec([], """():
+            class Foo:
+                def __length_hint__(self):
+                    1 / 0
+            return Foo()
+        """)
+        try:
+            assert length_hint(space, w_foo, 3)
+        except OperationError, e:
+            assert e.match(space, space.w_ZeroDivisionError)
+        else:
+            assert False, 'ZeroDivisionError expected'