[pypy-commit] pypy matrixmath: tests, implementation of flat

mattip noreply at buildbot.pypy.org
Sat Nov 26 19:58:20 CET 2011


Author: mattip
Branch: matrixmath
Changeset: r49830:db1df6669568
Date: 2011-11-26 20:47 +0200
http://bitbucket.org/pypy/pypy/changeset/db1df6669568/

Log:	tests, implementation of flat

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
@@ -312,6 +312,7 @@
     def get_offset(self):
         return 0
 
+
 class BaseArray(Wrappable):
     _attrs_ = ["invalidates", "signature", "shape", "strides", "backstrides",
                "start", 'order']
@@ -771,6 +772,9 @@
         return space.wrap(NDimSlice(self, new_sig, self.start, strides[:], 
                            backstrides[:], shape[:]))
 
+    def descr_get_flatiter(self, space):
+        return space.wrap(FlatIterator(self))
+
     def getitem(self, item):
         raise NotImplementedError
 
@@ -1177,6 +1181,7 @@
     size = GetSetProperty(BaseArray.descr_get_size),
 
     T = GetSetProperty(BaseArray.descr_get_transpose),
+    flat = GetSetProperty(BaseArray.descr_get_flatiter),
 
     mean = interp2app(BaseArray.descr_mean),
     sum = interp2app(BaseArray.descr_sum),
@@ -1191,3 +1196,33 @@
 
     copy = interp2app(BaseArray.descr_copy),
 )
+
+def descr_new_flatiter(space, w_object):
+    assert isinstance(w_object,BaseArray)
+    i = FlatIterator(w_object)
+    return i
+
+
+class FlatIterator(Wrappable):
+    _attrs_ = ["next"]
+
+    _immutable_fields_ = ['shapelen', ]
+
+    def __init__(self, arr):
+        self.arr = arr.get_concrete()
+        self.iter = arr.start_iter()
+        self.shapelen = len(arr.shape)
+
+    def descr_next(self, space):
+        if self.iter.done():
+            raise OperationError(space.w_StopIteration,space.wrap(''))
+        retVal = self.arr.eval(self.iter)
+        self.iter = self.iter.next(self.shapelen)
+        return retVal.wrap(space)
+            
+
+FlatIterator.typedef = TypeDef(
+    'flatiter',
+    __new__ = interp2app(descr_new_flatiter),
+    next = interp2app(FlatIterator.descr_next),
+)
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
@@ -947,6 +947,16 @@
         b[:, 0, 0] = 1000
         assert(a[0, 0, :] == [1000, 1000, 1000]).all()
 
+    def test_flatiter(self):
+        from numpypy import array
+        a = array([[10,30],[40,60]])
+        f_iter = a.flat
+        assert f_iter.next() == 10
+        assert f_iter.next() == 30
+        assert f_iter.next() == 40
+        assert f_iter.next() == 60
+        raises(StopIteration, "f_iter.next()")
+
 
 class AppTestSupport(object):
     def setup_class(cls):


More information about the pypy-commit mailing list