implement NPY_ARRAY_SAME_KIND_CASTING and use in np.take

Author: JuliaPoo

numpy/_core/src/multiarray/arrayobject.h

@@ -55,6 +55,14 @@ static const int NPY_ARRAY_WAS_PYTHON_COMPLEX = (1 << 28);
 static const int NPY_ARRAY_WAS_INT_AND_REPLACED = (1 << 27);
 static const int NPY_ARRAY_WAS_PYTHON_LITERAL = (1 << 30 | 1 << 29 | 1 << 28);
 
+/*
+ * This flag allows same kind casting, similar to NPY_ARRAY_FORCECAST.
+ *
+ * An array never has this flag set; they're only used as parameter
+ * flags to the various FromAny functions.
+ */
+static const int NPY_ARRAY_SAME_KIND_CASTING = (1 << 26);
+
 #ifdef __cplusplus
 }
 #endif

numpy/_core/src/multiarray/ctors.c

@@ -7,6 +7,7 @@
 #include <structmember.h>
 
 #include "numpy/arrayobject.h"
+#include "arrayobject.h"
 #include "numpy/arrayscalars.h"
 
 #include "numpy/npy_math.h"
@@ -1908,6 +1909,10 @@ PyArray_FromArray(PyArrayObject *arr, PyArray_Descr *newtype, int flags)
         newtype->elsize = oldtype->elsize;
     }
 
+    if (flags & NPY_ARRAY_SAME_KIND_CASTING) {
+        casting = NPY_SAME_KIND_CASTING;
+    }
+
     /* If the casting if forced, use the 'unsafe' casting rule */
     if (flags & NPY_ARRAY_FORCECAST) {
         casting = NPY_UNSAFE_CASTING;

numpy/_core/src/multiarray/item_selection.c

@@ -244,23 +244,11 @@ PyArray_TakeFrom(PyArrayObject *self0, PyObject *indices0, int axis,
         return NULL;
     }
 
-#if NPY_SIZEOF_INTP == NPY_SIZEOF_INT
-    PyArrayObject *tmp;
-    tmp = (PyArrayObject *)PyArray_ContiguousFromAny(indices0,
-                                                     NPY_INT64,
-                                                     0, 0);
-    if (tmp == NULL) {
-        goto fail;
-    }
-
-    indices = (PyArrayObject *)PyArray_Cast(tmp, NPY_INTP);
-    Py_DECREF(tmp);
-#else
-    indices = (PyArrayObject *)PyArray_ContiguousFromAny(indices0,
-                                                         NPY_INT64,
-                                                         0, 0);
-#endif
-
+    indices = (PyArrayObject *)PyArray_FromAny(indices0,
+                                               PyArray_DescrFromType(NPY_INTP),
+                                               0, 0,
+                                               NPY_ARRAY_SAME_KIND_CASTING,
+                                               NULL);
     if (indices == NULL) {
         goto fail;
     }