Add keyword argument to

numpy · ngoldbaum · Feb 29, 2024 · Nov 17, 2023 · Dec 8, 2023 · Dec 15, 2023
commit 8f5d130f8d44b9f9a0effe4ba727c5920622071d
diff --git a/numpy/_core/_add_newdocs.py b/numpy/_core/_add_newdocs.py
@@ -2943,11 +2943,15 @@
 
 add_newdoc('numpy._core.multiarray', 'ndarray', ('__array__',
     """
-    a.__array__([dtype], /)
-
-    Returns either a new reference to self if dtype is not given or a new array
-    of provided data type if dtype is different from the current dtype of the
-    array.
+    a.__array__([dtype], /, *, copy=None)
+
+    For ``dtype`` parameter it returns either a new reference to self if
+    ``dtype`` is not given or a new array of provided data type if ``dtype``
+    is different from the current data type of the array.
+    For ``copy`` parameter it returns a new reference to self if
+    ``copy=False`` or ``copy=None`` and copying isn't enforced by ``dtype``
+    parameter. The method returns a new array for ``copy=True``, regardless of
+    ``dtype`` parameter.
 
     """))
 

diff --git a/numpy/_core/src/multiarray/methods.c b/numpy/_core/src/multiarray/methods.c
@@ -924,13 +924,16 @@ array_wraparray(PyArrayObject *self, PyObject *args)
 
 
 static PyObject *
-array_getarray(PyArrayObject *self, PyObject *args)
+array_getarray(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
     PyArray_Descr *newtype = NULL;
+    NPY_COPYMODE copy = NPY_COPY_IF_NEEDED;
+    static char *kwlist[] = {"", "copy", NULL};
     PyObject *ret;
 
-    if (!PyArg_ParseTuple(args, "|O&:__array__",
-                            PyArray_DescrConverter, &newtype)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O&$O&:__array__", kwlist,
+                                     PyArray_DescrConverter, &newtype,
+                                     PyArray_CopyConverter, &copy)) {
         Py_XDECREF(newtype);
         return NULL;
     }
@@ -960,13 +963,26 @@ array_getarray(PyArrayObject *self, PyObject *args)
         Py_INCREF(self);
     }
 
-    if ((newtype == NULL) || PyArray_EquivTypes(PyArray_DESCR(self), newtype)) {
-        return (PyObject *)self;
-    }
-    else {
+    if (copy == NPY_COPY_ALWAYS) {
+        if (newtype == NULL) {
+            newtype = PyArray_DESCR(self);
+        }
         ret = PyArray_CastToType(self, newtype, 0);
         Py_DECREF(self);
         return ret;
+    } else { // copy == NPY_COPY_IF_NEEDED || copy == NPY_COPY_NEVER
+        if (newtype == NULL || PyArray_EquivTypes(PyArray_DESCR(self), newtype)) {
+            return (PyObject *)self;
+        }
+        if (copy == NPY_COPY_IF_NEEDED) {
+            ret = PyArray_CastToType(self, newtype, 0);
+            Py_DECREF(self);
+            return ret;
+        } else { // copy == NPY_COPY_NEVER
+            PyErr_SetString(PyExc_ValueError,
+                            "Unable to avoid copy while creating an array from given array.");
+            return NULL;
+        }
     }
 }
 
@@ -2848,7 +2864,7 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = {
     /* for subtypes */
     {"__array__",
         (PyCFunction)array_getarray,
-        METH_VARARGS, NULL},
+        METH_VARARGS | METH_KEYWORDS, NULL},
     {"__array_finalize__",
         (PyCFunction)array_finalizearray,
         METH_O, NULL},

diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py
@@ -8456,6 +8456,19 @@ def __array__(self):
             with pytest.raises(ValueError):
                 np.array(arr, copy=copy)
 
+    def test___array__copy_arg(self):
+        a = np.ones((10, 10), dtype=int)
+
+        assert np.shares_memory(a, a.__array__())
+        assert not np.shares_memory(a, a.__array__(float))
+        assert not np.shares_memory(a, a.__array__(float, copy=None))
+        assert not np.shares_memory(a, a.__array__(copy=True))
+        assert np.shares_memory(a, a.__array__(copy=None))
+        assert np.shares_memory(a, a.__array__(copy=False))
+        assert np.shares_memory(a, a.__array__(int, copy=False))
+        with pytest.raises(ValueError):
+            np.shares_memory(a, a.__array__(float, copy=False))
+
     @pytest.mark.parametrize(
             "arr", [np.ones(()), np.arange(81).reshape((9, 9))])
     @pytest.mark.parametrize("order1", ["C", "F", None])