[pypy-commit] pypy default: Implement numpypy.ndarray.flatten

[alex_gaynor]

Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: 
Changeset: r51592:6eb004a15a12
Date: 2012-01-21 11:02 -0600
http://bitbucket.org/pypy/pypy/changeset/6eb004a15a12/

Log:	Implement numpypy.ndarray.flatten

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
@@ -430,9 +430,15 @@
     def descr_copy(self, space):
         return self.copy(space)
 
+    def descr_flatten(self, space):
+        return self.flatten(space)
+
     def copy(self, space):
         return self.get_concrete().copy(space)
 
+    def flatten(self, space):
+        return self.get_concrete().flatten(space)
+
     def descr_len(self, space):
         if len(self.shape):
             return space.wrap(self.shape[0])
@@ -769,6 +775,11 @@
     def copy(self, space):
         return Scalar(self.dtype, self.value)
 
+    def flatten(self, space):
+        array = W_NDimArray(self.size, [self.size], self.dtype)
+        array.setitem(0, self.value)
+        return array
+
     def fill(self, space, w_value):
         self.value = self.dtype.coerce(space, w_value)
 
@@ -1150,6 +1161,15 @@
         array.setslice(space, self)
         return array
 
+    def flatten(self, space):
+        array = W_NDimArray(self.size, [self.size], self.dtype, self.order)
+        if self.supports_fast_slicing():
+            array._fast_setslice(space, self)
+        else:
+            arr = SliceArray(array.shape, array.dtype, array, self)
+            array._sliceloop(arr)
+        return array
+
     def fill(self, space, w_value):
         self.setslice(space, scalar_w(space, self.dtype, w_value))
 
@@ -1379,6 +1399,7 @@
     fill = interp2app(BaseArray.descr_fill),
 
     copy = interp2app(BaseArray.descr_copy),
+    flatten = interp2app(BaseArray.descr_flatten),
     reshape = interp2app(BaseArray.descr_reshape),
     tolist = interp2app(BaseArray.descr_tolist),
 )
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
@@ -1039,6 +1039,20 @@
         a = array([5.0])
         assert a.std() == 0.0
 
+    def test_flatten(self):
+        from _numpypy import array
+
+        a = array([[1, 2, 3], [4, 5, 6]])
+        assert (a.flatten() == [1, 2, 3, 4, 5, 6]).all()
+        a = array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+        assert (a.flatten() == [1, 2, 3, 4, 5, 6, 7, 8]).all()
+        a = array([1, 2, 3, 4, 5, 6, 7, 8])
+        assert (a[::2].flatten() == [1, 3, 5, 7]).all()
+        a = array([1, 2, 3])
+        assert ((a + a).flatten() == [2, 4, 6]).all()
+        a = array(2)
+        assert (a.flatten() == [2]).all()
+
 
 class AppTestMultiDim(BaseNumpyAppTest):
     def test_init(self):