[pypy-commit] pypy default: test, fix converting object dtype to str

[mattip]

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

Log:	test, fix converting object dtype to str

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
@@ -604,10 +604,14 @@
             raise oefmt(space.w_NotImplementedError,
                         "astype(%s) not implemented yet",
                         new_dtype.get_name())
-        if new_dtype.num == NPY.STRING and new_dtype.elsize == 0:
-            if cur_dtype.num == NPY.STRING:
-                new_dtype = descriptor.variable_dtype(
-                    space, 'S' + str(cur_dtype.elsize))
+        if new_dtype.is_str() and new_dtype.elsize == 0:
+            elsize = 0
+            itype = cur_dtype.itemtype
+            for i in range(self.get_size()):
+                elsize = max(elsize, len(itype.str_format(self.implementation.getitem(i), add_quotes=False)))
+            new_dtype = descriptor.variable_dtype(
+                    space, 'S' + str(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",
diff --git a/pypy/module/micronumpy/test/test_object_arrays.py b/pypy/module/micronumpy/test/test_object_arrays.py
--- a/pypy/module/micronumpy/test/test_object_arrays.py
+++ b/pypy/module/micronumpy/test/test_object_arrays.py
@@ -83,8 +83,8 @@
     def test_complex_op(self):
         import numpy as np
         import sys
-        a = np.array(['abc', 'def'], dtype=object) 
-        b = np.array([1, 2, 3], dtype=object) 
+        a = np.array(['abc', 'def'], dtype=object)
+        b = np.array([1, 2, 3], dtype=object)
         c = np.array([complex(1, 1), complex(1, -1)], dtype=object)
         for arg in (a,b,c):
             assert (arg == np.real(arg)).all()
@@ -164,3 +164,11 @@
         a = np.array([(1, 'object')], dt)
         # Wrong way - should complain about writing buffer to object dtype
         raises(ValueError, np.array, [1, 'object'], dt)
+
+    def test_astype(self):
+        import numpy as np
+        a = np.array([b'a' * 100], dtype='O')
+        assert 'a' * 100 in str(a)
+        b = a.astype('S')
+        assert 'a' * 100 in str(b)
+