[pypy-commit] pypy default: __iter__ for flatiter

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r49881:e847fa0fc602
Date: 2011-11-28 10:46 +0200
http://bitbucket.org/pypy/pypy/changeset/e847fa0fc602/

Log:	__iter__ for flatiter

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
@@ -1225,8 +1225,12 @@
         self.iter = self.iter.next(self.shapelen)
         return result.wrap(space)
 
+    def descr_iter(self):
+        return self
+
 
 W_FlatIterator.typedef = TypeDef(
     'flatiter',
     next = interp2app(W_FlatIterator.descr_next),
+    __iter__ = interp2app(W_FlatIterator.descr_iter),
 )
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
@@ -963,7 +963,10 @@
         assert f_iter.next() == 60
         raises(StopIteration, "f_iter.next()")
         raises(TypeError, "flatiter()")
-
+        s = 0
+        for k in a.flat:
+            s += k
+        assert s == 140
 
 class AppTestSupport(object):
     def setup_class(cls):