diff options
| -rw-r--r-- | numpy/core/src/multiarray/array_coercion.c | 29 |
1 files changed, 17 insertions, 12 deletions
diff --git a/numpy/core/src/multiarray/array_coercion.c b/numpy/core/src/multiarray/array_coercion.c index 562e4f008..1559f3485 100644 --- a/numpy/core/src/multiarray/array_coercion.c +++ b/numpy/core/src/multiarray/array_coercion.c @@ -15,6 +15,7 @@ #include "common_dtype.h" #include "dtypemeta.h" +#include "abstractdtypes.h" #include "array_coercion.h" #include "ctors.h" #include "common.h" @@ -204,7 +205,7 @@ _PyArray_MapPyTypeToDType( * Lookup the DType for a registered known python scalar type. * * @param pytype Python Type to look up - * @return DType, None if it a known non-scalar, or NULL if an unknown object. + * @return DType, None if it is a known non-scalar, or NULL if an unknown object. */ static NPY_INLINE PyArray_DTypeMeta * npy_discover_dtype_from_pytype(PyTypeObject *pytype) @@ -212,21 +213,25 @@ npy_discover_dtype_from_pytype(PyTypeObject *pytype) PyObject *DType; if (pytype == &PyArray_Type) { - Py_INCREF(Py_None); - return (PyArray_DTypeMeta *)Py_None; + DType = Py_None; } - - DType = PyDict_GetItem(_global_pytype_to_type_dict, (PyObject *)pytype); - if (DType == NULL) { - /* the python type is not known */ - return NULL; + else if (pytype == &PyFloat_Type) { + DType = (PyObject *)&PyArray_PyFloatAbstractDType; + } + else if (pytype == &PyLong_Type) { + DType = (PyObject *)&PyArray_PyIntAbstractDType; } + else { + DType = PyDict_GetItem(_global_pytype_to_type_dict, + (PyObject *)pytype); - Py_INCREF(DType); - if (DType == Py_None) { - return (PyArray_DTypeMeta *)Py_None; + if (DType == NULL) { + /* the python type is not known */ + return NULL; + } } - assert(PyObject_TypeCheck(DType, (PyTypeObject *)&PyArrayDTypeMeta_Type)); + Py_INCREF(DType); + assert(DType == Py_None || PyObject_TypeCheck(DType, (PyTypeObject *)&PyArrayDTypeMeta_Type)); return (PyArray_DTypeMeta *)DType; } |