[pypy-commit] pypy default: Begin implementing axes parameter in np.transpose()

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: 
Changeset: r78489:7fb7e3c2065c
Date: 2015-07-07 17:45 +0100
http://bitbucket.org/pypy/pypy/changeset/7fb7e3c2065c/

Log:	Begin implementing axes parameter in np.transpose()

diff --git a/pypy/module/micronumpy/concrete.py b/pypy/module/micronumpy/concrete.py
--- a/pypy/module/micronumpy/concrete.py
+++ b/pypy/module/micronumpy/concrete.py
@@ -268,13 +268,15 @@
             view = chunks.apply(space, orig_arr)
             view.implementation.setslice(space, w_value)
 
-    def transpose(self, orig_array):
+    def transpose(self, orig_array, axes=None):
         if len(self.get_shape()) < 2:
             return self
         strides = []
         backstrides = []
         shape = []
-        for i in range(len(self.get_shape()) - 1, -1, -1):
+        if axes is None:
+            axes = range(len(self.get_shape()) - 1, -1, -1)
+        for i in axes:
             strides.append(self.get_strides()[i])
             backstrides.append(self.get_backstrides()[i])
             shape.append(self.get_shape()[i])
diff --git a/pypy/module/micronumpy/ndarray.py b/pypy/module/micronumpy/ndarray.py
--- a/pypy/module/micronumpy/ndarray.py
+++ b/pypy/module/micronumpy/ndarray.py
@@ -393,15 +393,24 @@
             w_shape = space.newtuple(args_w)
         return self.reshape(space, w_shape)
 
-    def descr_get_transpose(self, space):
-        return W_NDimArray(self.implementation.transpose(self))
+    def descr_get_transpose(self, space, axes=None):
+        return W_NDimArray(self.implementation.transpose(self, axes))
 
     def descr_transpose(self, space, args_w):
-        if not (len(args_w) == 0 or
+        if len(args_w) == 1 and space.isinstance_w(args_w[0], space.w_tuple):
+            args_w = space.fixedview(args_w[0])
+        if (len(args_w) == 0 or
                 len(args_w) == 1 and space.is_none(args_w[0])):
-            raise OperationError(space.w_NotImplementedError, space.wrap(
-                "axes unsupported for transpose"))
-        return self.descr_get_transpose(space)
+            return self.descr_get_transpose(space)
+        else:
+            axes = []
+            for w_arg in args_w:
+                try:
+                    axes.append(support.index_w(space, w_arg))
+                except OperationError:
+                    raise oefmt(space.w_TypeError, "an integer is required")
+            return self.descr_get_transpose(space, axes)
+
 
     @unwrap_spec(axis1=int, axis2=int)
     def descr_swapaxes(self, space, axis1, axis2):
diff --git a/pypy/module/micronumpy/test/test_ndarray.py b/pypy/module/micronumpy/test/test_ndarray.py
--- a/pypy/module/micronumpy/test/test_ndarray.py
+++ b/pypy/module/micronumpy/test/test_ndarray.py
@@ -2765,6 +2765,11 @@
         b = a.T
         assert b.shape == (3, 2, 4)
         assert(b[0, :, 0] == [0, 3]).all()
+
+        c = a.transpose((1, 0, 2))
+        assert c.shape == (2, 4, 3)
+        assert (c.transpose(1, 0, 2) == a).all()
+
         b[:, 0, 0] = 1000
         assert(a[0, 0, :] == [1000, 1000, 1000]).all()
         a = array(range(5))
@@ -2775,9 +2780,6 @@
         assert(b[:, 0] == a[0, :]).all()
         assert (a.transpose() == b).all()
         assert (a.transpose(None) == b).all()
-        import sys
-        if '__pypy__' in sys.builtin_module_names:
-            raises(NotImplementedError, a.transpose, (1, 0, 2))
 
     def test_flatiter(self):
         from numpy import array, flatiter, arange, zeros