[pypy-svn] r58656 - in pypy/branch/2.5-merge/pypy/module/itertools: . test

[arigo at codespeak.net]

Author: arigo
Date: Mon Oct  6 15:54:45 2008
New Revision: 58656

Modified:
   pypy/branch/2.5-merge/pypy/module/itertools/interp_itertools.py
   pypy/branch/2.5-merge/pypy/module/itertools/test/test_itertools.py
Log:
(iko, arigo)
Add a repr to 'count' and 'repeat'.


Modified: pypy/branch/2.5-merge/pypy/module/itertools/interp_itertools.py
==============================================================================
--- pypy/branch/2.5-merge/pypy/module/itertools/interp_itertools.py	(original)
+++ pypy/branch/2.5-merge/pypy/module/itertools/interp_itertools.py	Mon Oct  6 15:54:45 2008
@@ -23,6 +23,10 @@
 
         return self.space.wrap(c)
 
+    def repr_w(self):
+        s = 'count(%d)' % (self.c,)
+        return self.space.wrap(s)
+
 
 def W_Count___new__(space, w_subtype, firstval=0):
     return space.wrap(W_Count(space, firstval))
@@ -32,6 +36,7 @@
         __new__ = interp2app(W_Count___new__, unwrap_spec=[ObjSpace, W_Root, int]),
         __iter__ = interp2app(W_Count.iter_w, unwrap_spec=['self']),
         next = interp2app(W_Count.next_w, unwrap_spec=['self']),
+        __repr__ = interp2app(W_Count.repr_w, unwrap_spec=['self']),
         __doc__ = """Make an iterator that returns consecutive integers starting
     with n.  If not specified n defaults to zero. Does not currently
     support python long integers. Often used as an argument to imap()
@@ -72,6 +77,14 @@
     def iter_w(self):
         return self.space.wrap(self)
 
+    def repr_w(self):
+        objrepr = self.space.str_w(self.space.repr(self.w_obj))
+        if self.counting:
+            s = 'repeat(%s, %d)' % (objrepr, self.count)
+        else:
+            s = 'repeat(%s)' % (objrepr,)
+        return self.space.wrap(s)
+
 def W_Repeat___new__(space, w_subtype, w_obj, w_times=None):
     return space.wrap(W_Repeat(space, w_obj, w_times))
 
@@ -80,6 +93,7 @@
         __new__  = interp2app(W_Repeat___new__, unwrap_spec=[ObjSpace, W_Root, W_Root, W_Root]),
         __iter__ = interp2app(W_Repeat.iter_w, unwrap_spec=['self']),
         next     = interp2app(W_Repeat.next_w, unwrap_spec=['self']),
+        __repr__ = interp2app(W_Repeat.repr_w, unwrap_spec=['self']),
         __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

Modified: pypy/branch/2.5-merge/pypy/module/itertools/test/test_itertools.py
==============================================================================
--- pypy/branch/2.5-merge/pypy/module/itertools/test/test_itertools.py	(original)
+++ pypy/branch/2.5-merge/pypy/module/itertools/test/test_itertools.py	Mon Oct  6 15:54:45 2008
@@ -29,6 +29,14 @@
 
         raises(OverflowError, itertools.count, sys.maxint + 1)
 
+    def test_count_repr(self):
+        import itertools
+
+        it = itertools.count(123)
+        assert repr(it) == 'count(123)'
+        it.next()
+        assert repr(it) == 'count(124)'
+
     def test_repeat(self):
         import itertools
 
@@ -66,6 +74,21 @@
 
         raises(OverflowError, itertools.repeat, None, sys.maxint + 1)
 
+    def test_repeat_repr(self):
+        import itertools
+
+        it = itertools.repeat('foobar')
+        assert repr(it) == "repeat('foobar')"
+        it.next()
+        assert repr(it) == "repeat('foobar')"
+
+        it = itertools.repeat('foobar', 10)
+        assert repr(it) == "repeat('foobar', 10)"
+        it.next()
+        assert repr(it) == "repeat('foobar', 9)"
+        list(it)
+        assert repr(it) == "repeat('foobar', 0)"
+
     def test_takewhile(self):
         import itertools