Merge pull request #20106 from charris/backport-20088

BUG: core: result_type(0, np.timedelta64(4)) would seg. fault.
numpy · Oct 12, 2021 · c36c4e3 · c36c4e3
2 parents 707f650 + b30f189
commit c36c4e3
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 20 deletions.
diff --git a/numpy/core/src/multiarray/convert_datatype.c b/numpy/core/src/multiarray/convert_datatype.c
@@ -1536,6 +1536,40 @@ should_use_min_scalar(npy_intp narrs, PyArrayObject **arr,
 }
 
 
+/*
+ * Utility function used only in PyArray_ResultType for value-based logic.
+ * See that function for the meaning and contents of the parameters.
+ */
+static PyArray_Descr *
+get_descr_from_cast_or_value(
+        npy_intp i,
+        PyArrayObject *arrs[],
+        npy_intp ndtypes,
+        PyArray_Descr *descriptor,
+        PyArray_DTypeMeta *common_dtype)
+{
+    PyArray_Descr *curr;
+    if (NPY_LIKELY(i < ndtypes ||
+            !(PyArray_FLAGS(arrs[i-ndtypes]) & _NPY_ARRAY_WAS_PYSCALAR))) {
+        curr = PyArray_CastDescrToDType(descriptor, common_dtype);
+    }
+    else {
+        /*
+         * Unlike `PyArray_CastToDTypeAndPromoteDescriptors`, deal with
+         * plain Python values "graciously". This recovers the original
+         * value the long route, but it should almost never happen...
+         */
+        PyObject *tmp = PyArray_GETITEM(arrs[i-ndtypes],
+                                        PyArray_BYTES(arrs[i-ndtypes]));
+        if (tmp == NULL) {
+            return NULL;
+        }
+        curr = common_dtype->discover_descr_from_pyobject(common_dtype, tmp);
+        Py_DECREF(tmp);
+    }
+    return curr;
+}
+
 /*NUMPY_API
  *
  * Produces the result type of a bunch of inputs, using the same rules
@@ -1672,28 +1706,15 @@ PyArray_ResultType(
         result = common_dtype->default_descr(common_dtype);
     }
     else {
-        result = PyArray_CastDescrToDType(all_descriptors[0], common_dtype);
+        result = get_descr_from_cast_or_value(
+                    0, arrs, ndtypes, all_descriptors[0], common_dtype);
+        if (result == NULL) {
+            goto error;
+        }
 
         for (npy_intp i = 1; i < ndtypes+narrs; i++) {
-            PyArray_Descr *curr;
-            if (NPY_LIKELY(i < ndtypes ||
-                    !(PyArray_FLAGS(arrs[i-ndtypes]) & _NPY_ARRAY_WAS_PYSCALAR))) {
-                curr = PyArray_CastDescrToDType(all_descriptors[i], common_dtype);
-            }
-            else {
-                /*
-                 * Unlike `PyArray_CastToDTypeAndPromoteDescriptors` deal with
-                 * plain Python values "graciously". This recovers the original
-                 * value the long route, but it should almost never happen...
-                 */
-                PyObject *tmp = PyArray_GETITEM(
-                        arrs[i-ndtypes], PyArray_BYTES(arrs[i-ndtypes]));
-                if (tmp == NULL) {
-                    goto error;
-                }
-                curr = common_dtype->discover_descr_from_pyobject(common_dtype, tmp);
-                Py_DECREF(tmp);
-            }
+            PyArray_Descr *curr = get_descr_from_cast_or_value(
+                    i, arrs, ndtypes, all_descriptors[i], common_dtype);
             if (curr == NULL) {
                 goto error;
             }

diff --git a/numpy/core/tests/test_dtype.py b/numpy/core/tests/test_dtype.py
@@ -1539,3 +1539,11 @@ class mytype:
             # Tests that a dtype must have its type field set up to np.dtype
             # or in this case a builtin instance.
             create_custom_field_dtype(blueprint, mytype, 2)
+
+
+def test_result_type_integers_and_unitless_timedelta64():
+    # Regression test for gh-20077.  The following call of `result_type`
+    # would cause a seg. fault.
+    td = np.timedelta64(4)
+    result = np.result_type(0, td)
+    assert_dtype_equal(result, td.dtype)