[pypy-commit] pypy speedup-unpackiterable: add itertools.repeat.__len__

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: speedup-unpackiterable
Changeset: r56049:ddb50e25e3e9
Date: 2012-07-13 00:59 +0200
http://bitbucket.org/pypy/pypy/changeset/ddb50e25e3e9/

Log:	add itertools.repeat.__len__

diff --git a/pypy/module/itertools/interp_itertools.py b/pypy/module/itertools/interp_itertools.py
--- a/pypy/module/itertools/interp_itertools.py
+++ b/pypy/module/itertools/interp_itertools.py
@@ -112,6 +112,11 @@
             s = 'repeat(%s)' % (objrepr,)
         return self.space.wrap(s)
 
+    def len(self, space):
+        if self.count == -1 or not self.counting:
+            raise OperationError(space.w_TypeError, space.wrap('len() of unsized object'))
+        return space.wrap(self.count)
+
 def W_Repeat___new__(space, w_subtype, w_object, w_times=None):
     r = space.allocate_instance(W_Repeat, w_subtype)
     r.__init__(space, w_object, w_times)
@@ -124,6 +129,7 @@
         __iter__ = interp2app(W_Repeat.iter_w),
         next     = interp2app(W_Repeat.next_w),
         __repr__ = interp2app(W_Repeat.repr_w),
+        __len__  = interp2app(W_Repeat.len),
         __doc__  = """Make an iterator that returns object over and over again.
     Runs indefinitely unless the times argument is specified.  Used
     as argument to imap() for invariant parameters to the called
diff --git a/pypy/module/itertools/test/test_itertools.py b/pypy/module/itertools/test/test_itertools.py
--- a/pypy/module/itertools/test/test_itertools.py
+++ b/pypy/module/itertools/test/test_itertools.py
@@ -88,6 +88,14 @@
         list(it)
         assert repr(it) == "repeat('foobar', 0)"
 
+    def test_repeat_len(self):
+        import itertools
+
+        r = itertools.repeat('a', 15)
+        r.next()
+        assert len(r) == 14
+        raises(TypeError, "len(itertools.repeat('xkcd'))")
+
     def test_takewhile(self):
         import itertools