[pypy-commit] pypy default: test, implement kwargs to astype

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: 
Changeset: r77928:b49bb002e037
Date: 2015-06-05 18:16 +0300
http://bitbucket.org/pypy/pypy/changeset/b49bb002e037/

Log:	test, implement kwargs to astype

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
@@ -17,7 +17,6 @@
 from rpython.rtyper.annlowlevel import cast_gcref_to_instance
 from pypy.interpreter.baseobjspace import W_Root
 
-
 class BaseConcreteArray(object):
     _immutable_fields_ = ['dtype?', 'storage', 'start', 'size', 'shape[*]',
                           'strides[*]', 'backstrides[*]', 'order', 'gcstruct',
@@ -334,7 +333,7 @@
     def get_buffer(self, space, readonly):
         return ArrayBuffer(self, readonly)
 
-    def astype(self, space, dtype):
+    def astype(self, space, dtype, order):
         # copy the general pattern of the strides
         # but make the array storage contiguous in memory
         shape = self.get_shape()
@@ -350,7 +349,7 @@
         else:
             t_strides = []
             backstrides = []
-        impl = ConcreteArray(shape, dtype, self.order, t_strides, backstrides)
+        impl = ConcreteArray(shape, dtype, order, t_strides, backstrides)
         loop.setslice(space, impl.get_shape(), impl, self)
         return impl
 
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
@@ -20,6 +20,7 @@
 from pypy.module.micronumpy.flagsobj import W_FlagsObject
 from pypy.module.micronumpy.strides import get_shape_from_iterable, \
     shape_agreement, shape_agreement_multiple, is_c_contiguous, is_f_contiguous
+from pypy.module.micronumpy.casting import can_cast_array
 
 
 def _match_dot_shapes(space, left, right):
@@ -43,6 +44,15 @@
         raise oefmt(space.w_ValueError, "objects are not aligned")
     return out_shape, right_critical_dim
 
+def get_order(proto_order, order):
+    if order == 'C':
+        return 'C'
+    elif order == 'F':
+        return 'F'
+    elif order == 'K':
+        return proto_order
+    elif order == 'A':
+        return proto_order
 
 class __extend__(W_NDimArray):
     @jit.unroll_safe
@@ -592,10 +602,11 @@
         if self.is_scalar():
             return space.wrap(0)
         dtype = self.get_dtype().descr_newbyteorder(space, NPY.NATIVE)
-        contig = self.implementation.astype(space, dtype)
+        contig = self.implementation.astype(space, dtype, self.get_order())
         return contig.argsort(space, w_axis)
 
-    def descr_astype(self, space, w_dtype):
+    @unwrap_spec(order=str, casting=str, subok=bool, copy=bool)
+    def descr_astype(self, space, w_dtype, order='K', casting='unsafe', subok=True, copy=True):
         cur_dtype = self.get_dtype()
         new_dtype = space.interp_w(descriptor.W_Dtype, space.call_function(
             space.gettypefor(descriptor.W_Dtype), w_dtype))
@@ -607,9 +618,24 @@
             if cur_dtype.num == NPY.STRING:
                 new_dtype = descriptor.variable_dtype(
                     space, 'S' + str(cur_dtype.elsize))
+        if not can_cast_array(space, self, new_dtype, casting):
+            raise oefmt(space.w_TypeError, "Cannot cast array from %s to %s"
+                        "according to the rule %s", 
+                        space.str_w(self.get_dtype().descr_repr(space)),
+                        space.str_w(new_dtype.descr_repr(space)), casting)
+        order  = get_order(self.get_order(), order)
+        if (not copy and new_dtype == self.get_dtype() and order == self.get_order()
+                and (subok or type(self) is W_NDimArray)):
+            return self
         impl = self.implementation
-        new_impl = impl.astype(space, new_dtype)
-        return wrap_impl(space, space.type(self), self, new_impl)
+        new_impl = impl.astype(space, new_dtype, order)
+        if new_impl is None:
+            return self
+        if subok:
+            w_type = space.type(self)
+        else:
+            w_type = None
+        return wrap_impl(space, w_type, self, new_impl)
 
     def descr_get_base(self, space):
         impl = self.implementation
diff --git a/pypy/module/micronumpy/nditer.py b/pypy/module/micronumpy/nditer.py
--- a/pypy/module/micronumpy/nditer.py
+++ b/pypy/module/micronumpy/nditer.py
@@ -470,7 +470,7 @@
                                     "Iterator operand required copying or "
                                     "buffering for operand %d", i)
                     impl = self.seq[i].implementation
-                    new_impl = impl.astype(space, selfd)
+                    new_impl = impl.astype(space, selfd, self.order)
                     self.seq[i] = W_NDimArray(new_impl)
         else:
             #copy them from seq
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
@@ -2241,6 +2241,15 @@
         assert c.shape == b.shape
         assert c.strides == (8,)
 
+        exc = raises(TypeError, a.astype, 'i8', casting='safe')
+        assert exc.value.message.startswith(
+                "Cannot cast array from dtype('complex128') to dtype('int64')")
+        a = arange(6, dtype='f4').reshape(2, 3)
+        b = a.astype('f4', copy=False)
+        assert a is b
+        b = a.astype('f4', order='C', copy=False)
+        assert a is b
+
     def test_base(self):
         from numpy import array
         assert array(1).base is None