Merge pull request numpy#19173 from czgdp1807/never_copy

ENH: Add support for copy modes to NumPy

Author: mattip

numpy/__init__.py

@@ -110,7 +110,8 @@
 import warnings
 
 from ._globals import (
-    ModuleDeprecationWarning, VisibleDeprecationWarning, _NoValue
+    ModuleDeprecationWarning, VisibleDeprecationWarning,
+    _NoValue, _CopyMode
 )
 
 # We first need to detect if we're being called as part of the numpy setup

numpy/__init__.pyi

@@ -5,6 +5,7 @@ import mmap
 import ctypes as ct
 import array as _array
 import datetime as dt
+import enum
 from abc import abstractmethod
 from types import TracebackType, MappingProxyType
 from contextlib import ContextDecorator
@@ -1745,7 +1746,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
         order: _OrderKACF = ...,
         casting: _CastingKind = ...,
         subok: bool = ...,
-        copy: bool = ...,
+        copy: bool | _CopyMode = ...,
     ) -> NDArray[_ScalarType]: ...
     @overload
     def astype(
@@ -1754,7 +1755,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
         order: _OrderKACF = ...,
         casting: _CastingKind = ...,
         subok: bool = ...,
-        copy: bool = ...,
+        copy: bool | _CopyMode = ...,
     ) -> NDArray[Any]: ...
 
     @overload
@@ -2495,7 +2496,7 @@ class generic(_ArrayOrScalarCommon):
         order: _OrderKACF = ...,
         casting: _CastingKind = ...,
         subok: bool = ...,
-        copy: bool = ...,
+        copy: bool | _CopyMode = ...,
     ) -> _ScalarType: ...
     @overload
     def astype(
@@ -2504,7 +2505,7 @@ class generic(_ArrayOrScalarCommon):
         order: _OrderKACF = ...,
         casting: _CastingKind = ...,
         subok: bool = ...,
-        copy: bool = ...,
+        copy: bool | _CopyMode = ...,
     ) -> Any: ...
 
     # NOTE: `view` will perform a 0D->scalar cast,
@@ -3253,6 +3254,11 @@ trunc: _UFunc_Nin1_Nout1[L['trunc'], L[7], None]
 
 abs = absolute
 
+class _CopyMode(enum.Enum):
+    ALWAYS: L[True]
+    IF_NEEDED: L[False]
+    NEVER: L[2]
+
 # Warnings
 class ModuleDeprecationWarning(DeprecationWarning): ...
 class VisibleDeprecationWarning(UserWarning): ...

numpy/_globals.py

@@ -15,8 +15,11 @@ def foo(arg=np._NoValue):
 motivated this module.
 
 """
+import enum
+
 __ALL__ = [
-    'ModuleDeprecationWarning', 'VisibleDeprecationWarning', '_NoValue'
+    'ModuleDeprecationWarning', 'VisibleDeprecationWarning',
+    '_NoValue', '_CopyMode'
     ]
 
 
@@ -89,3 +92,38 @@ def __repr__(self):
 
 
 _NoValue = _NoValueType()
+
+
+class _CopyMode(enum.Enum):
+    """
+    An enumeration for the copy modes supported
+    by numpy.copy() and numpy.array(). The following three modes are supported,
+
+    - ALWAYS: This means that a deep copy of the input
+              array will always be taken.
+    - IF_NEEDED: This means that a deep copy of the input
+                 array will be taken only if necessary.
+    - NEVER: This means that the deep copy will never be taken.
+             If a copy cannot be avoided then a `ValueError` will be
+             raised.
+
+    Note that the buffer-protocol could in theory do copies.  NumPy currently
+    assumes an object exporting the buffer protocol will never do this.
+    """
+
+    ALWAYS = True
+    IF_NEEDED = False
+    NEVER = 2
+
+    def __bool__(self):
+        # For backwards compatiblity
+        if self == _CopyMode.ALWAYS:
+            return True
+
+        if self == _CopyMode.IF_NEEDED:
+            return False
+
+        raise ValueError(f"{self} is neither True nor False.")
+
+
+_CopyMode.__module__ = 'numpy'

numpy/array_api/_creation_functions.py

@@ -41,7 +41,7 @@ def asarray(
     *,
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
-    copy: Optional[bool] = None,
+    copy: Optional[Union[bool, np._CopyMode]] = None,
 ) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.asarray <numpy.asarray>`.
@@ -55,13 +55,13 @@ def asarray(
     _check_valid_dtype(dtype)
     if device not in ["cpu", None]:
         raise ValueError(f"Unsupported device {device!r}")
-    if copy is False:
+    if copy in (False, np._CopyMode.IF_NEEDED):
         # Note: copy=False is not yet implemented in np.asarray
         raise NotImplementedError("copy=False is not yet implemented")
     if isinstance(obj, Array):
         if dtype is not None and obj.dtype != dtype:
             copy = True
-        if copy is True:
+        if copy in (True, np._CopyMode.ALWAYS):
             return Array._new(np.array(obj._array, copy=True, dtype=dtype))
         return obj
     if dtype is None and isinstance(obj, int) and (obj > 2 ** 64 or obj < -(2 ** 63)):

numpy/array_api/tests/test_creation_functions.py

@@ -43,12 +43,18 @@ def test_asarray_copy():
     a[0] = 0
     assert all(b[0] == 1)
     assert all(a[0] == 0)
-    # Once copy=False is implemented, replace this with
-    # a = asarray([1])
-    # b = asarray(a, copy=False)
-    # a[0] = 0
-    # assert all(b[0] == 0)
+    a = asarray([1])
+    b = asarray(a, copy=np._CopyMode.ALWAYS)
+    a[0] = 0
+    assert all(b[0] == 1)
+    assert all(a[0] == 0)
+    a = asarray([1])
+    b = asarray(a, copy=np._CopyMode.NEVER)
+    a[0] = 0
+    assert all(b[0] == 0)
     assert_raises(NotImplementedError, lambda: asarray(a, copy=False))
+    assert_raises(NotImplementedError,
+                  lambda: asarray(a, copy=np._CopyMode.IF_NEEDED))
 
 
 def test_arange_errors():

numpy/core/include/numpy/ndarraytypes.h

@@ -868,7 +868,7 @@ typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *);
 
 /*
  * Always copy the array. Returned arrays are always CONTIGUOUS,
- * ALIGNED, and WRITEABLE.
+ * ALIGNED, and WRITEABLE. See also: NPY_ARRAY_ENSURENOCOPY = 0x4000.
  *
  * This flag may be requested in constructor functions.
  */
@@ -937,6 +937,13 @@ typedef int (PyArray_FinalizeFunc)(PyArrayObject *, PyObject *);
 #define NPY_ARRAY_UPDATEIFCOPY    0x1000 /* Deprecated in 1.14 */
 #define NPY_ARRAY_WRITEBACKIFCOPY 0x2000
 
+/*
+ * No copy may be made while converting from an object/array (result is a view)
+ *
+ * This flag may be requested in constructor functions.
+ */
+#define NPY_ARRAY_ENSURENOCOPY 0x4000
+
 /*
  * NOTE: there are also internal flags defined in multiarray/arrayobject.h,
  * which start at bit 31 and work down.

numpy/core/multiarray.pyi

@@ -50,6 +50,7 @@ from numpy import (
     _ModeKind,
     _SupportsBuffer,
     _IOProtocol,
+    _CopyMode,
     _NDIterFlagsKind,
     _NDIterOpFlagsKind,
 )
@@ -177,7 +178,7 @@ def array(
     object: _ArrayType,
     dtype: None = ...,
     *,
-    copy: bool = ...,
+    copy: bool | _CopyMode = ...,
     order: _OrderKACF = ...,
     subok: L[True],
     ndmin: int = ...,
@@ -188,7 +189,7 @@ def array(
     object: _ArrayLike[_SCT],
     dtype: None = ...,
     *,
-    copy: bool = ...,
+    copy: bool | _CopyMode = ...,
     order: _OrderKACF = ...,
     subok: bool = ...,
     ndmin: int = ...,
@@ -199,7 +200,7 @@ def array(
     object: object,
     dtype: None = ...,
     *,
-    copy: bool = ...,
+    copy: bool | _CopyMode = ...,
     order: _OrderKACF = ...,
     subok: bool = ...,
     ndmin: int = ...,
@@ -210,7 +211,7 @@ def array(
     object: Any,
     dtype: _DTypeLike[_SCT],
     *,
-    copy: bool = ...,
+    copy: bool | _CopyMode = ...,
     order: _OrderKACF = ...,
     subok: bool = ...,
     ndmin: int = ...,
@@ -221,7 +222,7 @@ def array(
     object: Any,
     dtype: DTypeLike,
     *,
-    copy: bool = ...,
+    copy: bool | _CopyMode = ...,
     order: _OrderKACF = ...,
     subok: bool = ...,
     ndmin: int = ...,

numpy/core/src/multiarray/_multiarray_tests.c.src

@@ -2363,6 +2363,17 @@ run_intp_converter(PyObject* NPY_UNUSED(self), PyObject *args)
     return tup;
 }
 
+/* used to test NPY_ARRAY_ENSURENOCOPY raises ValueError */
+static PyObject*
+npy_ensurenocopy(PyObject* NPY_UNUSED(self), PyObject* args)
+{
+    int flags = NPY_ARRAY_ENSURENOCOPY;
+    if (!PyArray_CheckFromAny(args, NULL, 0, 0, flags, NULL)) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef Multiarray_TestsMethods[] = {
     {"argparse_example_function",
          (PyCFunction)argparse_example_function,
@@ -2424,6 +2435,9 @@ static PyMethodDef Multiarray_TestsMethods[] = {
     {"npy_discard",
         npy_discard,
         METH_O, NULL},
+    {"npy_ensurenocopy",
+        npy_ensurenocopy,
+        METH_O, NULL},
     {"get_buffer_info",
         get_buffer_info,
         METH_VARARGS, NULL},

numpy/core/src/multiarray/array_coercion.c

@@ -858,6 +858,7 @@ PyArray_AdaptDescriptorToArray(PyArrayObject *arr, PyObject *dtype)
  *        (Initially it is a pointer to the user-provided head pointer).
  * @param fixed_DType User provided fixed DType class
  * @param flags Discovery flags (reporting and behaviour flags, see def.)
+ * @param never_copy Specifies if a copy is allowed during array creation.
  * @return The updated number of maximum dimensions (i.e. scalars will set
  *         this to the current dimensions).
  */
@@ -866,7 +867,8 @@ PyArray_DiscoverDTypeAndShape_Recursive(
         PyObject *obj, int curr_dims, int max_dims, PyArray_Descr**out_descr,
         npy_intp out_shape[NPY_MAXDIMS],
         coercion_cache_obj ***coercion_cache_tail_ptr,
-        PyArray_DTypeMeta *fixed_DType, enum _dtype_discovery_flags *flags)
+        PyArray_DTypeMeta *fixed_DType, enum _dtype_discovery_flags *flags,
+        int never_copy)
 {
     PyArrayObject *arr = NULL;
     PyObject *seq;
@@ -924,7 +926,7 @@ PyArray_DiscoverDTypeAndShape_Recursive(
             requested_descr = *out_descr;
         }
         arr = (PyArrayObject *)_array_from_array_like(obj,
-                requested_descr, 0, NULL);
+                requested_descr, 0, NULL, never_copy);
         if (arr == NULL) {
             return -1;
         }
@@ -1118,7 +1120,7 @@ PyArray_DiscoverDTypeAndShape_Recursive(
         max_dims = PyArray_DiscoverDTypeAndShape_Recursive(
                 objects[i], curr_dims + 1, max_dims,
                 out_descr, out_shape, coercion_cache_tail_ptr, fixed_DType,
-                flags);
+                flags, never_copy);
 
         if (max_dims < 0) {
             return -1;
@@ -1158,6 +1160,7 @@ PyArray_DiscoverDTypeAndShape_Recursive(
  *        The result may be unchanged (remain NULL) when converting a
  *        sequence with no elements. In this case it is callers responsibility
  *        to choose a default.
+ * @param never_copy Specifies that a copy is not allowed.
  * @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.
@@ -1170,7 +1173,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)
+        PyArray_Descr **out_descr, int never_copy)
 {
     coercion_cache_obj **coercion_cache_head = coercion_cache;
     *coercion_cache = NULL;
@@ -1215,7 +1218,7 @@ PyArray_DiscoverDTypeAndShape(
 
     int ndim = PyArray_DiscoverDTypeAndShape_Recursive(
             obj, 0, max_dims, out_descr, out_shape, &coercion_cache,
-            fixed_DType, &flags);
+            fixed_DType, &flags, never_copy);
     if (ndim < 0) {
         goto fail;
     }
@@ -1500,7 +1503,7 @@ _discover_array_parameters(PyObject *NPY_UNUSED(self),
     int ndim = PyArray_DiscoverDTypeAndShape(
             obj, NPY_MAXDIMS, shape,
             &coercion_cache,
-            fixed_DType, fixed_descriptor, (PyArray_Descr **)&out_dtype);
+            fixed_DType, fixed_descriptor, (PyArray_Descr **)&out_dtype, 0);
     Py_XDECREF(fixed_DType);
     Py_XDECREF(fixed_descriptor);
     if (ndim < 0) {

numpy/core/src/multiarray/array_coercion.h

@@ -31,7 +31,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);
+        PyArray_Descr **out_descr, int never_copy);
 
 NPY_NO_EXPORT int
 PyArray_ExtractDTypeAndDescriptor(PyObject *dtype,