Merge pull request numpy#26215 from mtsokol/one-copy-__array__

API: Enforce one copy for ``__array__`` when ``copy=True``

Author: mtsokol

numpy/_core/src/multiarray/array_coercion.c

@@ -99,6 +99,7 @@ enum _dtype_discovery_flags {
     DISCOVER_TUPLES_AS_ELEMENTS = 1 << 4,
     MAX_DIMS_WAS_REACHED = 1 << 5,
     DESCRIPTOR_WAS_SET = 1 << 6,
+    COPY_WAS_CREATED_BY__ARRAY__ = 1 << 7,
 };
 
 
@@ -1027,15 +1028,19 @@ PyArray_DiscoverDTypeAndShape_Recursive(
             /* __array__ may be passed the requested descriptor if provided */
             requested_descr = *out_descr;
         }
+        int was_copied_by__array__ = 0;
         arr = (PyArrayObject *)_array_from_array_like(obj,
-                requested_descr, 0, NULL, copy);
+                requested_descr, 0, NULL, copy, &was_copied_by__array__);
         if (arr == NULL) {
             return -1;
         }
         else if (arr == (PyArrayObject *)Py_NotImplemented) {
             Py_DECREF(arr);
             arr = NULL;
         }
+        if (was_copied_by__array__ == 1) {
+            *flags |= COPY_WAS_CREATED_BY__ARRAY__;
+        }
     }
     if (arr != NULL) {
         /*
@@ -1170,6 +1175,15 @@ PyArray_DiscoverDTypeAndShape_Recursive(
         return -1;
     }
 
+    /*
+     * For a sequence we need to make a copy of the final aggreate anyway.
+     * There's no need to pass explicit `copy=True`, so we switch
+     * to `copy=None` (copy if needed).
+     */
+    if (copy == 1) {
+        copy = -1;
+    }
+
     /* Recursive call for each sequence item */
     for (Py_ssize_t i = 0; i < size; i++) {
         max_dims = PyArray_DiscoverDTypeAndShape_Recursive(
@@ -1217,6 +1231,8 @@ PyArray_DiscoverDTypeAndShape_Recursive(
  *        to choose a default.
  * @param copy Specifies the copy behavior. -1 is corresponds to copy=None,
  *        0 to copy=False, and 1 to copy=True in the Python API.
+ * @param was_copied_by__array__ Set to 1 if it can be assumed that a copy was
+ *        made by implementor.
  * @return dimensions of the discovered object or -1 on error.
  *         WARNING: If (and only if) the output is a single array, the ndim
  *         returned _can_ exceed the maximum allowed number of dimensions.
@@ -1229,7 +1245,7 @@ PyArray_DiscoverDTypeAndShape(
         npy_intp out_shape[NPY_MAXDIMS],
         coercion_cache_obj **coercion_cache,
         PyArray_DTypeMeta *fixed_DType, PyArray_Descr *requested_descr,
-        PyArray_Descr **out_descr, int copy)
+        PyArray_Descr **out_descr, int copy, int *was_copied_by__array__)
 {
     coercion_cache_obj **coercion_cache_head = coercion_cache;
     *coercion_cache = NULL;
@@ -1282,6 +1298,10 @@ PyArray_DiscoverDTypeAndShape(
         goto fail;
     }
 
+    if (was_copied_by__array__ != NULL && flags & COPY_WAS_CREATED_BY__ARRAY__) {
+        *was_copied_by__array__ = 1;
+    }
+
     if (NPY_UNLIKELY(flags & FOUND_RAGGED_ARRAY)) {
         /*
          * If max-dims was reached and the dimensions reduced, this is ragged.
@@ -1396,7 +1416,7 @@ _discover_array_parameters(PyObject *NPY_UNUSED(self),
     int ndim = PyArray_DiscoverDTypeAndShape(
             obj, NPY_MAXDIMS, shape,
             &coercion_cache,
-            dt_info.dtype, dt_info.descr, (PyArray_Descr **)&out_dtype, 0);
+            dt_info.dtype, dt_info.descr, (PyArray_Descr **)&out_dtype, 0, NULL);
     Py_XDECREF(dt_info.dtype);
     Py_XDECREF(dt_info.descr);
     if (ndim < 0) {

numpy/_core/src/multiarray/array_coercion.h

@@ -40,7 +40,7 @@ PyArray_DiscoverDTypeAndShape(
         npy_intp out_shape[NPY_MAXDIMS],
         coercion_cache_obj **coercion_cache,
         PyArray_DTypeMeta *fixed_DType, PyArray_Descr *requested_descr,
-        PyArray_Descr **out_descr, int copy);
+        PyArray_Descr **out_descr, int copy, int *was_copied_by__array__);
 
 NPY_NO_EXPORT PyObject *
 _discover_array_parameters(PyObject *NPY_UNUSED(self),

numpy/_core/src/multiarray/arrayobject.c

@@ -251,7 +251,7 @@ PyArray_CopyObject(PyArrayObject *dest, PyObject *src_object)
      */
     ndim = PyArray_DiscoverDTypeAndShape(src_object,
             PyArray_NDIM(dest), dims, &cache,
-            NPY_DTYPE(PyArray_DESCR(dest)), PyArray_DESCR(dest), &dtype, 1);
+            NPY_DTYPE(PyArray_DESCR(dest)), PyArray_DESCR(dest), &dtype, 1, NULL);
     if (ndim < 0) {
         return -1;
     }

numpy/_core/src/multiarray/common.c

@@ -119,7 +119,7 @@ PyArray_DTypeFromObject(PyObject *obj, int maxdims, PyArray_Descr **out_dtype)
     int ndim;
 
     ndim = PyArray_DiscoverDTypeAndShape(
-            obj, maxdims, shape, &cache, NULL, NULL, out_dtype, 1);
+            obj, maxdims, shape, &cache, NULL, NULL, out_dtype, 1, NULL);
     if (ndim < 0) {
         return -1;
     }

numpy/_core/src/multiarray/ctors.c

@@ -1429,6 +1429,8 @@ _array_from_buffer_3118(PyObject *memoryview)
  * @param writeable whether the result must be writeable.
  * @param context Unused parameter, must be NULL (should be removed later).
  * @param copy Specifies the copy behavior.
+ * @param was_copied_by__array__ Set to 1 if it can be assumed that a copy
+ *        was made by implementor.
  *
  * @returns The array object, Py_NotImplemented if op is not array-like,
  *          or NULL with an error set. (A new reference to Py_NotImplemented
@@ -1437,7 +1439,7 @@ _array_from_buffer_3118(PyObject *memoryview)
 NPY_NO_EXPORT PyObject *
 _array_from_array_like(PyObject *op,
         PyArray_Descr *requested_dtype, npy_bool writeable, PyObject *context,
-        int copy) {
+        int copy, int *was_copied_by__array__) {
     PyObject* tmp;
 
     /*
@@ -1485,7 +1487,8 @@ _array_from_array_like(PyObject *op,
     }
 
     if (tmp == Py_NotImplemented) {
-        tmp = PyArray_FromArrayAttr_int(op, requested_dtype, copy);
+        tmp = PyArray_FromArrayAttr_int(
+                op, requested_dtype, copy, was_copied_by__array__);
         if (tmp == NULL) {
             return NULL;
         }
@@ -1572,13 +1575,17 @@ PyArray_FromAny_int(PyObject *op, PyArray_Descr *in_descr,
 
     // Default is copy = None
     int copy = -1;
+    int was_copied_by__array__ = 0;
 
     if (flags & NPY_ARRAY_ENSURENOCOPY) {
         copy = 0;
+    } else if (flags & NPY_ARRAY_ENSURECOPY) {
+        copy = 1;
     }
 
     ndim = PyArray_DiscoverDTypeAndShape(
-            op, NPY_MAXDIMS, dims, &cache, in_DType, in_descr, &dtype, copy);
+            op, NPY_MAXDIMS, dims, &cache, in_DType, in_descr, &dtype,
+            copy, &was_copied_by__array__);
 
     if (ndim < 0) {
         return NULL;
@@ -1615,6 +1622,9 @@ PyArray_FromAny_int(PyObject *op, PyArray_Descr *in_descr,
         assert(cache->converted_obj == op);
         arr = (PyArrayObject *)(cache->arr_or_sequence);
         /* we may need to cast or assert flags (e.g. copy) */
+        if (was_copied_by__array__ == 1) {
+            flags = flags & ~NPY_ARRAY_ENSURECOPY;
+        }
         PyObject *res = PyArray_FromArray(arr, dtype, flags);
         npy_unlink_coercion_cache(cache);
         return res;
@@ -1937,7 +1947,7 @@ PyArray_FromArray(PyArrayObject *arr, PyArray_Descr *newtype, int flags)
     }
 
     if (copy) {
-        if (flags & NPY_ARRAY_ENSURENOCOPY ) {
+        if (flags & NPY_ARRAY_ENSURENOCOPY) {
             PyErr_SetString(PyExc_ValueError, npy_no_copy_err_msg);
             Py_DECREF(newtype);
             return NULL;
@@ -2497,12 +2507,14 @@ check_or_clear_and_warn_error_if_due_to_copy_kwarg(PyObject *kwnames)
  *        NOTE: For copy == -1 it passes `op.__array__(copy=None)`,
  *              for copy == 0, `op.__array__(copy=False)`, and
  *              for copy == 1, `op.__array__(copy=True).
+ * @param was_copied_by__array__ Set to 1 if it can be assumed that a copy
+ *        was made by implementor.
  * @returns NotImplemented if `__array__` is not defined or a NumPy array
  *          (or subclass).  On error, return NULL.
  */
 NPY_NO_EXPORT PyObject *
-PyArray_FromArrayAttr_int(
-        PyObject *op, PyArray_Descr *descr, int copy)
+PyArray_FromArrayAttr_int(PyObject *op, PyArray_Descr *descr, int copy,
+                          int *was_copied_by__array__)
 {
     PyObject *new;
     PyObject *array_meth;
@@ -2589,10 +2601,13 @@ PyArray_FromArrayAttr_int(
         Py_DECREF(new);
         return NULL;
     }
-    if (must_copy_but_copy_kwarg_unimplemented) {
-        /* TODO: As of NumPy 2.0 this path is only reachable by C-API. */
-        Py_SETREF(new, PyArray_NewCopy((PyArrayObject *)new, NPY_KEEPORDER));
+    /* TODO: Remove was_copied_by__array__ argument */
+    if (was_copied_by__array__ != NULL && copy == 1 &&
+        must_copy_but_copy_kwarg_unimplemented == 0) {
+        /* We can assume that a copy was made */
+        *was_copied_by__array__ = 1;
     }
+
     return new;
 }
 
@@ -2607,7 +2622,7 @@ PyArray_FromArrayAttr(PyObject *op, PyArray_Descr *typecode, PyObject *context)
         return NULL;
     }
 
-    return PyArray_FromArrayAttr_int(op, typecode, 0);
+    return PyArray_FromArrayAttr_int(op, typecode, 0, NULL);
 }
 
 

numpy/_core/src/multiarray/ctors.h

@@ -54,7 +54,7 @@ PyArray_New(
 NPY_NO_EXPORT PyObject *
 _array_from_array_like(PyObject *op,
         PyArray_Descr *requested_dtype, npy_bool writeable, PyObject *context,
-        int copy);
+        int copy, int *was_copied_by__array__);
 
 NPY_NO_EXPORT PyObject *
 PyArray_FromAny_int(PyObject *op, PyArray_Descr *in_descr,
@@ -84,8 +84,8 @@ NPY_NO_EXPORT PyObject *
 PyArray_FromInterface(PyObject *input);
 
 NPY_NO_EXPORT PyObject *
-PyArray_FromArrayAttr_int(
-        PyObject *op, PyArray_Descr *descr, int copy);
+PyArray_FromArrayAttr_int(PyObject *op, PyArray_Descr *descr, int copy,
+                          int *was_copied_by__array__);
 
 NPY_NO_EXPORT PyObject *
 PyArray_FromArrayAttr(PyObject *op, PyArray_Descr *typecode,

numpy/_core/tests/test_array_coercion.py

@@ -54,7 +54,9 @@ def __init__(self, a):
             self.a = a
 
         def __array__(self, dtype=None, copy=None):
-            return self.a
+            if dtype is None:
+                return self.a
+            return self.a.astype(dtype)
 
     yield param(ArrayDunder, id="__array__")
 

numpy/_core/tests/test_multiarray.py

@@ -8452,10 +8452,9 @@ def __array__(self, dtype=None, copy=None):
         for copy in self.true_vals:
             res = np.array(arr, copy=copy)
             assert_array_equal(res, base_arr)
-            # An additional copy is currently forced by numpy in this case,
-            # you could argue, numpy does not trust the ArrayLike. This
-            # may be open for change:
-            assert res is not base_arr
+            # An additional copy is no longer forced by NumPy in this case.
+            # NumPy trusts the ArrayLike made a copy:
+            assert res is base_arr
 
         for copy in self.if_needed_vals + self.false_vals:
             res = np.array(arr, copy=copy)
@@ -8488,9 +8487,11 @@ def __array__(self, dtype=None):
         assert_array_equal(arr, base_arr)
         assert arr is base_arr
 
-        # As of NumPy 2, explicitly passing copy=True does not trigger passing
-        # it to __array__ (deprecation warning is not triggered).
-        arr = np.array(a, copy=True)
+        # As of NumPy 2.1, explicitly passing copy=True does trigger passing
+        # it to __array__ (deprecation warning is triggered).
+        with pytest.warns(DeprecationWarning,
+                          match="__array__.*must implement.*'copy'"):
+            arr = np.array(a, copy=True)
         assert_array_equal(arr, base_arr)
         assert arr is not base_arr
 
@@ -8501,10 +8502,45 @@ def __array__(self, dtype=None):
                     match=r"Unable to avoid copy(.|\n)*numpy_2_0_migration_guide.html"):
                 np.array(a, copy=False)
 
+    def test___array__copy_once(self):
+        size = 100
+        base_arr = np.zeros((size, size))
+        copy_arr = np.zeros((size, size))
+
+        class ArrayRandom:
+            def __init__(self):
+                self.true_passed = False
+
+            def __array__(self, dtype=None, copy=None):
+                if copy:
+                    self.true_passed = True
+                    return copy_arr
+                else:
+                    return base_arr
+
+        arr_random = ArrayRandom()
+        first_copy = np.array(arr_random, copy=True)
+        assert arr_random.true_passed
+        assert first_copy is copy_arr
+
+        arr_random = ArrayRandom()
+        no_copy = np.array(arr_random, copy=False)
+        assert not arr_random.true_passed
+        assert no_copy is base_arr
+
+        arr_random = ArrayRandom()
+        _ = np.array([arr_random], copy=True)
+        assert not arr_random.true_passed
+
+        arr_random = ArrayRandom()
+        second_copy = np.array(arr_random, copy=True, order="F")
+        assert arr_random.true_passed
+        assert second_copy is not copy_arr
+
     @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
     def test__array__reference_leak(self):
         class NotAnArray:
-            def __array__(self):
+            def __array__(self, dtype=None, copy=None):
                 raise NotImplementedError()
 
         x = NotAnArray()

numpy/_core/tests/test_protocols.py

@@ -35,8 +35,8 @@ def test_array_called():
     class Wrapper:
         val = '0' * 100
 
-        def __array__(self, result=None, copy=None):
-            return np.array([self.val], dtype=object)
+        def __array__(self, dtype=None, copy=None):
+            return np.array([self.val], dtype=dtype, copy=copy)
 
 
     wrapped = Wrapper()