MAINT: Put array tagging into a static inline function

this also effectively fixes some corner cases in np.result_type

3 files changed, 59 insertions, 41 deletions

diff --git a/numpy/core/src/multiarray/abstractdtypes.h b/numpy/core/src/multiarray/abstractdtypes.h
index 42c192cac..b3f547b64 100644
--- a/numpy/core/src/multiarray/abstractdtypes.h
+++ b/numpy/core/src/multiarray/abstractdtypes.h
@@ -1,6 +1,7 @@
 #ifndef NUMPY_CORE_SRC_MULTIARRAY_ABSTRACTDTYPES_H_
 #define NUMPY_CORE_SRC_MULTIARRAY_ABSTRACTDTYPES_H_
 
+#include "arrayobject.h"
 #include "dtypemeta.h"
 
 
@@ -16,4 +17,56 @@ NPY_NO_EXPORT extern PyArray_DTypeMeta PyArray_PyComplexAbstractDType;
 NPY_NO_EXPORT int
 initialize_and_map_pytypes_to_dtypes(void);
 
+
+/*
+ * When we get a Python int, float, or complex, we may have to use weak
+ * promotion logic.
+ * To implement this, we sometimes have to tag the converted (temporary)
+ * array when the original object was a Python scalar.
+ *
+ * @param obj The original Python object.
+ * @param arr The array into which the Python object was converted.
+ * @param[in,out] **dtype A pointer to the array's DType, if not NULL it will be
+ *        replaced with the abstract DType.
+ * @return 0 if the `obj` was not a python scalar, and 1 if it was.
+ */
+static NPY_INLINE int
+npy_mark_tmp_array_if_pyscalar(
+        PyObject *obj, PyArrayObject *arr, PyArray_DTypeMeta **dtype)
+{
+    /*
+     * We check the array dtype for two reasons: First, booleans are
+     * integer subclasses.  Second, an int, float, or complex could have
+     * a custom DType registered, and then we should use that.
+     * Further, `np.float64` is a double subclass, so must reject it.
+     */
+    if (PyLong_Check(obj)
+        && (PyArray_ISINTEGER(arr) || PyArray_ISOBJECT(arr))) {
+        ((PyArrayObject_fields *)arr)->flags |= NPY_ARRAY_WAS_PYTHON_INT;
+        if (dtype != NULL) {
+            Py_INCREF(&PyArray_PyIntAbstractDType);
+            Py_SETREF(*dtype, &PyArray_PyIntAbstractDType);
+        }
+        return 1;
+    }
+    else if (PyFloat_Check(obj) && !PyArray_IsScalar(obj, Double)
+             && PyArray_TYPE(arr) == NPY_DOUBLE) {
+        ((PyArrayObject_fields *)arr)->flags |= NPY_ARRAY_WAS_PYTHON_FLOAT;
+        if (dtype != NULL) {
+            Py_INCREF(&PyArray_PyFloatAbstractDType);
+            Py_SETREF(*dtype, &PyArray_PyFloatAbstractDType);
+        }
+        return 1;
+    }
+    else if (PyComplex_Check(obj) && PyArray_TYPE(arr) == NPY_CDOUBLE) {
+        ((PyArrayObject_fields *)arr)->flags |= NPY_ARRAY_WAS_PYTHON_COMPLEX;
+        if (dtype != NULL) {
+            Py_INCREF(&PyArray_PyComplexAbstractDType);
+            Py_SETREF(*dtype, &PyArray_PyComplexAbstractDType);
+        }
+        return 1;
+    }
+    return 0;
+}
+
 #endif  /* NUMPY_CORE_SRC_MULTIARRAY_ABSTRACTDTYPES_H_ */
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 10114ce8b..16069d619 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -3537,19 +3537,11 @@ array_result_type(PyObject *NPY_UNUSED(dummy), PyObject *args)
             if (arr[narr] == NULL) {
                 goto finish;
             }
-            // TODO: Relax the exact check here, see ufunc code!
-            if (PyLong_CheckExact(obj)) {
-                ((PyArrayObject_fields *)arr[narr])->flags |= (
-                        NPY_ARRAY_WAS_PYTHON_INT);
-            }
-            else if (PyFloat_CheckExact(obj)) {
-                ((PyArrayObject_fields *)arr[narr])->flags |= (
-                        NPY_ARRAY_WAS_PYTHON_FLOAT);
-            }
-            else if (PyComplex_CheckExact(obj)) {
-                ((PyArrayObject_fields *)arr[narr])->flags |= (
-                        NPY_ARRAY_WAS_PYTHON_COMPLEX);
-            }
+            /*
+             * Mark array if it was a python scalar (we do not need the actual
+             * DType here yet, this is figured out inside ResultType.
+             */
+            npy_mark_tmp_array_if_pyscalar(obj, arr[narr], NULL);
             ++narr;
         }
         else {
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 4d7b49b3b..b7e390abb 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -1008,35 +1008,8 @@ convert_ufunc_arguments(PyUFuncObject *ufunc,
          * this.  This is because the legacy dtype resolution makes use of
          * `np.can_cast(operand, dtype)`.  The flag is local to this use, but
          * necessary to propagate the information to the legacy type resolution.
-         *
-         * We check `out_op_DTypes` for two reasons: First, booleans are
-         * integer subclasses.  Second, an int, float, or complex could have
-         * a custom DType registered, and then we should use that.
          */
-        if (PyLong_Check(obj)
-                && (PyTypeNum_ISINTEGER(out_op_DTypes[i]->type_num)
-                    || out_op_DTypes[i]->type_num == NPY_OBJECT)) {
-            Py_INCREF(&PyArray_PyIntAbstractDType);
-            Py_SETREF(out_op_DTypes[i], &PyArray_PyIntAbstractDType);
-            ((PyArrayObject_fields *)out_op[i])->flags |= (
-                    NPY_ARRAY_WAS_PYTHON_INT);
-            *promoting_pyscalars = NPY_TRUE;
-        }
-        else if (PyFloat_Check(obj)
-                && out_op_DTypes[i]->type_num == NPY_DOUBLE
-                && !PyArray_IsScalar(obj, Double)) {
-            Py_INCREF(&PyArray_PyFloatAbstractDType);
-            Py_SETREF(out_op_DTypes[i], &PyArray_PyFloatAbstractDType);
-            ((PyArrayObject_fields *)out_op[i])->flags |= (
-                    NPY_ARRAY_WAS_PYTHON_FLOAT);
-            *promoting_pyscalars = NPY_TRUE;
-        }
-        else if (PyComplex_Check(obj)
-                && out_op_DTypes[i]->type_num == NPY_CDOUBLE) {
-            Py_INCREF(&PyArray_PyComplexAbstractDType);
-            Py_SETREF(out_op_DTypes[i], &PyArray_PyComplexAbstractDType);
-            ((PyArrayObject_fields *)out_op[i])->flags |= (
-                    NPY_ARRAY_WAS_PYTHON_COMPLEX);
+        if (npy_mark_tmp_array_if_pyscalar(obj, out_op[i], &out_op_DTypes[i])) {
             *promoting_pyscalars = NPY_TRUE;
         }
     }