improved handling of 2d arrays

Author: flexatone

src/methods.c

@@ -259,7 +259,6 @@ astype_array(PyObject* m, PyObject* args) {
         &dtype_spec)) {
         return NULL;
     }
-    // AK_CHECK_NUMPY_ARRAY(a);
     PyArrayObject* array = (PyArrayObject*)a;
 
     PyArray_Descr* dtype = NULL;
@@ -294,16 +293,25 @@ astype_array(PyObject* m, PyObject* args) {
                     return NULL;
                 }
                 PyObject** data = (PyObject**)PyArray_DATA((PyArrayObject*)result);
-                npy_intp size = PyArray_SIZE(array);
 
-                for (npy_intp i = 0; i < size; ++i) {
-                    PyObject* item = PyArray_Scalar(PyArray_GETPTR1(array, i), array_dt, a);
+                PyArrayIterObject* it = (PyArrayIterObject*)PyArray_IterNew(a);
+                if (!it) {
+                    Py_DECREF(result);
+                    return NULL;
+                }
+
+                npy_intp i = 0;
+                while (it->index < it->size) {
+                    PyObject* item = PyArray_ToScalar(it->dataptr, array);
                     if (!item) {
                         Py_DECREF(result);
+                        Py_DECREF(it);
                         return NULL;
                     }
-                    data[i] = item;
+                    data[i++] = item;
+                    PyArray_ITER_NEXT(it);
                 }
+                Py_DECREF(it);
                 PyArray_CLEARFLAGS((PyArrayObject *)result, NPY_ARRAY_WRITEABLE);
                 return result;
             }

test/test_astype_array.py

@@ -54,7 +54,40 @@ def test_astype_array_b3(self) -> None:
         a2 = astype_array(a1, np.object_)
         self.assertEqual(a2.dtype, np.dtype(np.object_))
         self.assertFalse(a2.flags.writeable)
-        import ipdb; ipdb.set_trace()
         self.assertEqual(
                 list(list(a) for a in a2),
                 [[np.datetime64('2021'), np.datetime64('2024')], [np.datetime64('1984'), np.datetime64('1642')]])
+
+    def test_astype_array_b4(self) -> None:
+        a1 = np.array(['2021', '2024', '1532', '1984', '1642', '899'], dtype=np.datetime64).reshape((2, 3))
+
+        a2 = astype_array(a1, np.object_)
+        self.assertEqual(a2.dtype, np.dtype(np.object_))
+        self.assertEqual(a2.shape, (2, 3))
+        self.assertFalse(a2.flags.writeable)
+        self.assertEqual(
+                list(list(a) for a in a2),
+                [[np.datetime64('2021'), np.datetime64('2024'), np.datetime64('1532')],
+                 [np.datetime64('1984'), np.datetime64('1642'), np.datetime64('899')]])
+
+    def test_astype_array_c(self) -> None:
+        with self.assertRaises(TypeError):
+            _ = astype_array([3, 4, 5], np.int64)
+
+
+    def test_astype_array_d1(self) -> None:
+        a1 = np.array([10, 20, 30], dtype=np.int64)
+        a2 = astype_array(a1)
+
+        self.assertEqual(a2.dtype, np.dtype(np.float64))
+        self.assertEqual(a2.shape, (3,))
+        self.assertFalse(a2.flags.writeable)
+
+
+    def test_astype_array_d2(self) -> None:
+        a1 = np.array([10, 20, 30], dtype=np.int64)
+        a2 = astype_array(a1, None)
+
+        self.assertEqual(a2.dtype, np.dtype(np.float64))
+        self.assertEqual(a2.shape, (3,))
+        self.assertFalse(a2.flags.writeable)