[pypy-commit] pypy default: added flatiter.__len__

alex_gaynor noreply at buildbot.pypy.org
Sat Feb 4 02:00:32 CET 2012


Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: 
Changeset: r52074:2610a9ed0d7c
Date: 2012-02-03 20:00 -0500
http://bitbucket.org/pypy/pypy/changeset/2610a9ed0d7c/

Log:	added flatiter.__len__

diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -796,7 +796,7 @@
                                self.left.create_sig(), self.right.create_sig())
 
 class ResultArray(Call2):
-    def __init__(self, child, size, shape, dtype, res=None, order='C'): 
+    def __init__(self, child, size, shape, dtype, res=None, order='C'):
         if res is None:
             res = W_NDimArray(size, shape, dtype, order)
         Call2.__init__(self, None, 'assign', shape, dtype, dtype, res, child)
@@ -824,7 +824,7 @@
             frame.next(len(self.right.shape))
         else:
             frame.cur_value = self.identity.convert_to(self.calc_dtype)
-    
+
     def create_sig(self):
         if self.name == 'logical_and':
             done_func = done_if_false
@@ -1321,6 +1321,9 @@
     def descr_iter(self):
         return self
 
+    def descr_len(self, space):
+        return space.wrap(self.size)
+
     def descr_index(self, space):
         return space.wrap(self.index)
 
@@ -1396,7 +1399,7 @@
         return signature.FlatSignature(self.base.dtype)
 
     def create_iter(self, transforms=None):
-        return ViewIterator(self.base.start, self.base.strides, 
+        return ViewIterator(self.base.start, self.base.strides,
                     self.base.backstrides,
                     self.base.shape).apply_transformations(self.base,
                                                            transforms)
@@ -1407,14 +1410,17 @@
 W_FlatIterator.typedef = TypeDef(
     'flatiter',
     __iter__ = interp2app(W_FlatIterator.descr_iter),
+    __len__ = interp2app(W_FlatIterator.descr_len),
     __getitem__ = interp2app(W_FlatIterator.descr_getitem),
     __setitem__ = interp2app(W_FlatIterator.descr_setitem),
+
     __eq__ = interp2app(BaseArray.descr_eq),
     __ne__ = interp2app(BaseArray.descr_ne),
     __lt__ = interp2app(BaseArray.descr_lt),
     __le__ = interp2app(BaseArray.descr_le),
     __gt__ = interp2app(BaseArray.descr_gt),
     __ge__ = interp2app(BaseArray.descr_ge),
+
     base = GetSetProperty(W_FlatIterator.descr_base),
     index = GetSetProperty(W_FlatIterator.descr_index),
     coords = GetSetProperty(W_FlatIterator.descr_coords),
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -1489,24 +1489,26 @@
     def test_flatiter_view(self):
         from _numpypy import arange
         a = arange(10).reshape(5, 2)
-        #no == yet.
-        # a[::2].flat == [0, 1, 4, 5, 8, 9]
-        isequal = True
-        for y,z in zip(a[::2].flat, [0, 1, 4, 5, 8, 9]):
-            if y != z:
-                isequal = False
-        assert isequal == True
+        assert (a[::2].flat == [0, 1, 4, 5, 8, 9]).all()
 
     def test_flatiter_transpose(self):
         from _numpypy import arange
-        a = arange(10).reshape(2,5).T
+        a = arange(10).reshape(2, 5).T
         b = a.flat
         assert (b[:5] == [0, 5, 1, 6, 2]).all()
         b.next()
         b.next()
         b.next()
         assert b.index == 3
-        assert b.coords == (1,1)
+        assert b.coords == (1, 1)
+
+    def test_flatiter_len(self):
+        from _numpypy import arange
+
+        assert len(arange(10).flat) == 10
+        assert len(arange(10).reshape(2, 5).flat) == 10
+        assert len(arange(10)[:2].flat) == 2
+        assert len((arange(2) + arange(2)).flat) == 2
 
     def test_slice_copy(self):
         from _numpypy import zeros


More information about the pypy-commit mailing list