[pypy-commit] pypy numpy-refactor: some flatiter support

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-refactor
Changeset: r57182:e3d43b2ad3ec
Date: 2012-09-06 17:27 +0200
http://bitbucket.org/pypy/pypy/changeset/e3d43b2ad3ec/

Log:	some flatiter support

diff --git a/pypy/module/micronumpy/interp_flatiter.py b/pypy/module/micronumpy/interp_flatiter.py
--- a/pypy/module/micronumpy/interp_flatiter.py
+++ b/pypy/module/micronumpy/interp_flatiter.py
@@ -1,10 +1,29 @@
 
 from pypy.interpreter.baseobjspace import Wrappable
-from pypy.interpreter.typedef import TypeDef
+from pypy.interpreter.error import OperationError
+from pypy.interpreter.typedef import TypeDef, interp2app
+from pypy.rlib import jit
 
 class W_FlatIterator(Wrappable):
-    pass
+    def __init__(self, arr):
+        self.arr = arr
+        self.iter = self.arr.create_iter()
+        self.index = 0
+
+    def descr_next(self, space):
+        if self.iter.done():
+            raise OperationError(space.w_StopIteration, space.w_None)
+        w_res = self.iter.getitem()
+        self.iter.next()
+        self.index += 1
+        return w_res
+
+    def descr_iter(self):
+        return self
 
 W_FlatIterator.typedef = TypeDef(
     'flatiter',
+    __iter__ = interp2app(W_FlatIterator.descr_iter),    
+
+    next = interp2app(W_FlatIterator.descr_next),
 )
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
@@ -3,9 +3,10 @@
 from pypy.interpreter.typedef import TypeDef, GetSetProperty
 from pypy.interpreter.gateway import interp2app, unwrap_spec
 from pypy.module.micronumpy.base import W_NDimArray, convert_to_array
-from pypy.module.micronumpy import interp_dtype, interp_ufuncs, support
+from pypy.module.micronumpy import interp_dtype, interp_ufuncs
 from pypy.module.micronumpy.strides import find_shape_and_elems,\
      get_shape_from_iterable
+from pypy.module.micronumpy.interp_flatiter import W_FlatIterator
 from pypy.module.micronumpy.interp_support import unwrap_axis_arg
 from pypy.module.micronumpy.appbridge import get_appbridge_cache
 from pypy.module.micronumpy import loop
@@ -217,6 +218,9 @@
     def descr_repeat(self, space, repeats, w_axis=None):
         return repeat(space, self, repeats, w_axis)
 
+    def descr_get_flatiter(self, space):
+        return space.wrap(W_FlatIterator(self))
+
     # --------------------- operations ----------------------------
 
     def _unaryop_impl(ufunc_name):
@@ -468,6 +472,7 @@
     ravel = interp2app(W_NDimArray.descr_ravel),
     repeat = interp2app(W_NDimArray.descr_repeat),
     swapaxes = interp2app(W_NDimArray.descr_swapaxes),
+    flat = GetSetProperty(W_NDimArray.descr_get_flatiter),
 )
 
 @unwrap_spec(ndmin=int, copy=bool, subok=bool)