ENH: iter: Switch the iterator to use PyArray_ResultType

3 files changed, 40 insertions, 163 deletions

diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index 0f58f66a2..db84046cb 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -1275,8 +1275,8 @@ add_newdoc('numpy.core.multiarray', 'min_scalar_type',
     min_scalar_type(a)
 
     For scalar ``a``, returns the data type with the smallest size
-    and smallest scalar kind which can hold its value.  For vector ``a``,
-    returns the vector's dtype unmodified.
+    and smallest scalar kind which can hold its value.  For non-scalar
+    array ``a``, returns the vector's dtype unmodified.
 
     As a special case, floating point values are not reduced to integers.
 
diff --git a/numpy/core/src/multiarray/convert_datatype.c b/numpy/core/src/multiarray/convert_datatype.c
index a13eea6b5..c22e6a0b7 100644
--- a/numpy/core/src/multiarray/convert_datatype.c
+++ b/numpy/core/src/multiarray/convert_datatype.c
@@ -1016,8 +1016,9 @@ PyArray_ResultType(npy_intp narrs, PyArrayObject **arr,
             PyArray_Descr *tmp = PyArray_DESCR(arr[i]);
             int tmp_is_small_unsigned = 0;
             /*
-             * If it's a scalar, find the min scalar type.  The function is expanded here so that
-             * we can flag whether we've got an unsigned integer which would fit an a signed integer
+             * If it's a scalar, find the min scalar type. The function
+             * is expanded here so that we can flag whether we've got an
+             * unsigned integer which would fit an a signed integer
              * of the same size, something not exposed in the public API.
              */
             if (PyArray_NDIM(arr[i]) == 0 && PyTypeNum_ISNUMBER(tmp->type_num)) {
@@ -1031,7 +1032,8 @@ PyArray_ResultType(npy_intp narrs, PyArrayObject **arr,
                 npy_clongdouble value;
 #endif
                 tmp->f->copyswap(&value, data, swap, NULL);
-                type_num = min_scalar_type_num((char *)&value, tmp->type_num, &tmp_is_small_unsigned);
+                type_num = min_scalar_type_num((char *)&value,
+                                        tmp->type_num, &tmp_is_small_unsigned);
                 tmp = PyArray_DescrFromType(type_num);
                 if (tmp == NULL) {
                     Py_XDECREF(ret);
@@ -1055,8 +1057,15 @@ PyArray_ResultType(npy_intp narrs, PyArrayObject **arr,
                 printf(" (%d) ", ret_is_small_unsigned);
                 printf("\n");
 #endif
-                tmpret = promote_types(tmp, ret, tmp_is_small_unsigned, ret_is_small_unsigned);
-                ret_is_small_unsigned = tmp_is_small_unsigned && ret_is_small_unsigned;
+                tmpret = promote_types(tmp, ret, tmp_is_small_unsigned,
+                                                    ret_is_small_unsigned);
+                if (tmpret == NULL) {
+                    Py_DECREF(tmp);
+                    Py_DECREF(ret);
+                    return NULL;
+                }
+                ret_is_small_unsigned = tmp_is_small_unsigned &&
+                                        ret_is_small_unsigned;
                 Py_DECREF(tmp);
                 Py_DECREF(ret);
                 ret = tmpret;
@@ -1073,6 +1082,11 @@ PyArray_ResultType(npy_intp narrs, PyArrayObject **arr,
             else {
                 if (ret_is_small_unsigned) {
                     tmpret = promote_types(tmp, ret, 0, ret_is_small_unsigned);
+                    if (tmpret == NULL) {
+                        Py_DECREF(tmp);
+                        Py_DECREF(ret);
+                        return NULL;
+                    }
                 }
                 else {
                     tmpret = PyArray_PromoteTypes(tmp, ret);
@@ -1083,6 +1097,11 @@ PyArray_ResultType(npy_intp narrs, PyArrayObject **arr,
         }
     }
 
+    if (ret == NULL) {
+        PyErr_SetString(PyExc_TypeError,
+                "no arrays or types available to calculate result type");
+    }
+
     return ret;
 }
 
diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
index 7d45d25d3..b7b928051 100644
--- a/numpy/core/src/multiarray/new_iterator.c.src
+++ b/numpy/core/src/multiarray/new_iterator.c.src
@@ -277,9 +277,8 @@ npyiter_shrink_ndim(NpyIter *iter, npy_intp new_ndim);
 static PyArray_Descr *
 npyiter_get_common_dtype(npy_intp niter, PyArrayObject **op,
                         char *op_itflags, PyArray_Descr **op_dtype,
+                        PyArray_Descr **op_request_dtypes,
                         int only_inputs, int output_scalars);
-static int
-npyiter_promote_types(int type1, int type2);
 
 static PyArrayObject *
 npyiter_new_temp_array(NpyIter *iter, PyTypeObject *subtype,
@@ -495,6 +494,7 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
 
         dtype = npyiter_get_common_dtype(niter, op,
                                     op_itflags, op_dtype,
+                                    op_request_dtypes,
                                     only_inputs,
                                     output_scalars);
         if (dtype == NULL) {
@@ -3935,173 +3935,31 @@ npyiter_get_priority_subtype(npy_intp niter, PyArrayObject **op,
 static PyArray_Descr *
 npyiter_get_common_dtype(npy_intp niter, PyArrayObject **op,
                         char *op_itflags, PyArray_Descr **op_dtype,
+                        PyArray_Descr **op_request_dtypes,
                         int only_inputs, int output_scalars)
 {
-    PyArray_Descr *scalar_dtype = NULL, *array_dtype = NULL;
-    int scalar_kind = NPY_NOSCALAR;
     npy_intp iiter;
-    /* 0 = don't use, 1 = scalar, 2 = array */
-    char op_category[NPY_MAXARGS];
+    npy_intp narrs = 0, ndtypes = 0;
+    PyArrayObject *arrs[NPY_MAXARGS];
+    PyArray_Descr *dtypes[NPY_MAXARGS];
 
-    /* Determine which operands are scalars and which dtypes to use */
     for (iiter = 0; iiter < niter; ++iiter) {
         if (op_dtype[iiter] != NULL &&
                     (!only_inputs || (op_itflags[iiter]&NPY_OP_ITFLAG_READ))) {
-            if ((op[iiter] == NULL && output_scalars) ||
-                    PyArray_NDIM(op[iiter]) == 0) {
-                op_category[iiter] = 1;
-            }
-            else {
-                op_category[iiter] = 2;
-            }
-        }
-        else {
-            op_category[iiter] = 0;
-        }
-    }
-
-    /* First promote all the scalar dtypes */
-    for (iiter = 0; iiter < niter; ++iiter) {
-        /* Process all the scalar inputs (category = 1) */
-        if (op_category[iiter] == 1) {
-            if (scalar_dtype == NULL) {
-                scalar_dtype = op_dtype[iiter];
-                Py_INCREF(scalar_dtype);
-                /*
-                 * Note: because dtype and op[iiter] may not have the same
-                 *       type, this value could be (slightly) wrong in
-                 *       very few cases.
-                 */
-                scalar_kind = PyArray_ScalarKind(scalar_dtype->type_num,
-                                                            &op[iiter]);
-            }
-            else {
-                PyArray_Descr *dtype = op_dtype[iiter];
-                int kind, typenum;
-
-                if (PyTypeNum_ISEXTENDED(scalar_dtype->type_num) ||
-                                PyTypeNum_ISEXTENDED(dtype->type_num)) {
-                    PyErr_SetString(PyExc_TypeError,
-                            "Iterator allocated output could not "
-                            "determine a data type based on the inputs");
-                    Py_DECREF(scalar_dtype);
-                    return NULL;
-                }
-                /*
-                 * Note: because dtype and op[iiter] may not have the same
-                 *       type, this value could be (slightly) wrong in
-                 *       very few cases.
-                 */
-                kind = PyArray_ScalarKind(dtype->type_num, &op[iiter]);
-                if (kind > scalar_kind) {
-                    scalar_kind = kind;
-                }
-                typenum = npyiter_promote_types(scalar_dtype->type_num,
-                                                  dtype->type_num);
-                if (typenum == NPY_NOTYPE) {
-                    PyErr_SetString(PyExc_TypeError,
-                            "Iterator allocated output could not "
-                            "determine a data type based on the inputs");
-                    Py_DECREF(scalar_dtype);
-                    return NULL;
-                }
-
-                dtype = PyArray_DescrFromType(typenum);
-                Py_DECREF(scalar_dtype);
-                scalar_dtype = dtype;
-            }
-        }
-    }
-
-    /* Then promote all the array dtypes */
-    for (iiter = 0; iiter < niter; ++iiter) {
-        /* Process all the array inputs (category = 2) */
-        if (op_category[iiter] == 2) {
-            if (array_dtype == NULL) {
-                array_dtype = op_dtype[iiter];
-                Py_INCREF(array_dtype);
+            /* If no dtype was requested and the op is a scalar, pass the op */
+            if ((op_request_dtypes == NULL ||
+                            op_request_dtypes[iiter] == NULL) &&
+                                            PyArray_NDIM(op[iiter]) == 0) {
+                arrs[narrs++] = op[iiter];
             }
+            /* Otherwise just pass in the dtype */
             else {
-                PyArray_Descr *dtype = op_dtype[iiter];
-                int typenum;
-
-                if (PyTypeNum_ISEXTENDED(array_dtype->type_num) ||
-                                PyTypeNum_ISEXTENDED(dtype->type_num)) {
-                    PyErr_SetString(PyExc_TypeError,
-                            "Iterator allocated output could not "
-                            "determine a data type based on the inputs");
-                    Py_XDECREF(scalar_dtype);
-                    return NULL;
-                }
-                typenum = npyiter_promote_types(array_dtype->type_num,
-                                                  dtype->type_num);
-                if (typenum == NPY_NOTYPE) {
-                    PyErr_SetString(PyExc_TypeError,
-                            "Iterator allocated output could not "
-                            "determine a data type based on the inputs");
-                    Py_XDECREF(scalar_dtype);
-                    return NULL;
-                }
-                            
-                dtype = PyArray_DescrFromType(typenum);
-                Py_DECREF(array_dtype);
-                array_dtype = dtype;
+                dtypes[ndtypes++] = op_dtype[iiter];
             }
         }
     }
 
-    /* Now combine the scalar and array types */
-    if (scalar_dtype != NULL && array_dtype != NULL) {
-        PyArray_Descr *dtype;
-        if (PyArray_CanCoerceScalar(scalar_dtype->type_num,
-                                    array_dtype->type_num,
-                                    scalar_kind)) {
-            dtype = PyArray_DescrFromType(array_dtype->type_num);
-            Py_DECREF(scalar_dtype);
-            Py_DECREF(array_dtype);
-            return dtype;
-        }
-        else {
-            int typenum;
-
-            typenum = npyiter_promote_types(scalar_dtype->type_num,
-                                            array_dtype->type_num);
-            Py_DECREF(scalar_dtype);
-            Py_DECREF(array_dtype);
-            return PyArray_DescrFromType(typenum);
-        }
-    }
-    else if (array_dtype != NULL) {
-        return array_dtype;
-    }
-    else if (scalar_dtype != NULL) {
-        return scalar_dtype;
-    }
-    else {
-        PyErr_SetString(PyExc_TypeError,
-                "Iterator allocated output could not "
-                "determine a data type based on the inputs");
-        return NULL;
-    }
-}
-
-/*
- * TODO This is a slow way to do it and depends on the type
- * number ordering.  Should be based on a table like scalar kinds.
- */
-static int
-npyiter_promote_types(int type1, int type2)
-{
-    int i;
-
-    for(i = 0; i < NPY_NTYPES; ++i) {
-        if (PyArray_CanCastSafely(type1, i) &&
-                            PyArray_CanCastSafely(type2, i)) {
-            return i;
-        }
-    }
-
-    return NPY_NOTYPE;
+    return PyArray_ResultType(narrs, arrs, ndtypes, dtypes);
 }
 
 static int