WIP: Implement weak scalar logic

numpy · Jun 15, 2022 · edb369d · edb369d
1 parent 7c14383
commit edb369d
Show file tree

Hide file tree

Showing 7 changed files with 171 additions and 189 deletions.
diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
@@ -872,17 +872,6 @@ typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *);
  */
 #define NPY_ARRAY_ENSUREARRAY     0x0040
 
-#if defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD
-    /*
-     * Dual use of the ENSUREARRAY flag, to indicate that this was converted
-     * from a python float, int, or complex.
-     * An array using this flag must be a temporary array that can never
-     * leave the C internals of NumPy.  Even if it does, ENSUREARRAY is
-     * absolutely safe to abuse, since it already is a base class array :).
-     */
-    #define _NPY_ARRAY_WAS_PYSCALAR   0x0040
-#endif  /* NPY_INTERNAL_BUILD */
-
 /*
  * Make sure that the strides are in units of the element size Needed
  * for some operations with record-arrays.

diff --git a/numpy/core/src/multiarray/array_assign_scalar.c b/numpy/core/src/multiarray/array_assign_scalar.c
@@ -243,8 +243,11 @@ PyArray_AssignRawScalar(PyArrayObject *dst,
     }
 
     /* Check the casting rule */
-    if (!can_cast_scalar_to(src_dtype, src_data,
-                            PyArray_DESCR(dst), casting)) {
+    // TODO: This is incorrect, we need the information if the input was a
+    //       Python scalar?!  Or rather, we want a "inspect" cast-safety
+    //       so to speak.  A special scalar assignment just for this purpose
+    //       _probably_ that should be handled earlier, though!
+    if (!PyArray_CanCastTypeTo(src_dtype, PyArray_DESCR(dst), casting)) {
         npy_set_invalid_cast_error(
                 src_dtype, PyArray_DESCR(dst), casting, NPY_TRUE);
         return -1;

diff --git a/numpy/core/src/multiarray/arrayobject.h b/numpy/core/src/multiarray/arrayobject.h
@@ -26,4 +26,20 @@ array_might_be_written(PyArrayObject *obj);
  */
 static const int NPY_ARRAY_WARN_ON_WRITE = (1 << 31);
 
+
+/*
+ * These flags are used internally to indicate an array that was previously
+ * a Python scalar (int, float, complex).  The dtype of such an array should
+ * be considered as any integer, floating, or complex rather than the explicit
+ * dtype attached to the array.
+ *
+ * These flags must only be used in local context when the array in question
+ * is not returned.  Use three flags, to avoid having to double check the
+ * actual dtype when the flags are used.
+ */
+static const int NPY_ARRAY_WAS_PYTHON_INT = (1 << 30);
+static const int NPY_ARRAY_WAS_PYTHON_FLOAT = (1 << 29);
+static const int NPY_ARRAY_WAS_PYTHON_COMPLEX = (1 << 28);
+static const int NPY_ARRAY_WAS_PYTHON_LITERAL = (1 << 30 | 1 << 29 | 1 << 28);
+
 #endif  /* NUMPY_CORE_SRC_MULTIARRAY_ARRAYOBJECT_H_ */