API: Add shape argument to numpy.reshape

mtsokol · mtsokol · commit e3bf8e5d836c · 2024-04-17T14:33:36.000+02:00
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -1713,11 +1713,19 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
 
     @overload
     def reshape(
-        self, shape: _ShapeLike, /, *, order: _OrderACF = ...
+        self,
+        shape: _ShapeLike,
+        /,
+        *,
+        order: _OrderACF = ...,
+        copy: None | bool = ...,
     ) -> ndarray[Any, _DType_co]: ...
     @overload
     def reshape(
-        self, *shape: SupportsIndex, order: _OrderACF = ...
+        self,
+        *shape: SupportsIndex,
+        order: _OrderACF = ...,
+        copy: None | bool = ...,
     ) -> ndarray[Any, _DType_co]: ...
 
     @overload
diff --git a/numpy/_core/_add_newdocs.py b/numpy/_core/_add_newdocs.py
@@ -3824,7 +3824,7 @@
 
 add_newdoc('numpy._core.multiarray', 'ndarray', ('reshape',
     """
-    a.reshape(shape, /, *, order='C')
+    a.reshape(shape, /, *, order='C', copy=None)
 
     Returns an array containing the same data with a new shape.
 
diff --git a/numpy/_core/fromnumeric.py b/numpy/_core/fromnumeric.py
@@ -206,13 +206,13 @@ def take(a, indices, axis=None, out=None, mode='raise'):
     return _wrapfunc(a, 'take', indices, axis=axis, out=out, mode=mode)
 
 
-def _reshape_dispatcher(a, newshape, order=None):
+def _reshape_dispatcher(a, /, newshape=None, shape=None, *,
+                        order=None, copy=None):
     return (a,)
 
 
-# not deprecated --- copy if necessary, view otherwise
 @array_function_dispatch(_reshape_dispatcher)
-def reshape(a, newshape, order='C'):
+def reshape(a, /, newshape=None, shape=None, *, order='C', copy=None):
     """
     Gives a new shape to an array without changing its data.
 
@@ -221,13 +221,16 @@ def reshape(a, newshape, order='C'):
     a : array_like
         Array to be reshaped.
     newshape : int or tuple of ints
+        Replaced by ``shape`` argument. Retained for backward
+        compatibility.
+    shape : int or tuple of ints
         The new shape should be compatible with the original shape. If
         an integer, then the result will be a 1-D array of that length.
         One shape dimension can be -1. In this case, the value is
         inferred from the length of the array and remaining dimensions.
     order : {'C', 'F', 'A'}, optional
-        Read the elements of `a` using this index order, and place the
-        elements into the reshaped array using this index order.  'C'
+        Read the elements of ``a`` using this index order, and place the
+        elements into the reshaped array using this index order. 'C'
         means to read / write the elements using C-like index order,
         with the last axis index changing fastest, back to the first
         axis index changing slowest. 'F' means to read / write the
@@ -236,8 +239,12 @@ def reshape(a, newshape, order='C'):
         the 'C' and 'F' options take no account of the memory layout of
         the underlying array, and only refer to the order of indexing.
         'A' means to read / write the elements in Fortran-like index
-        order if `a` is Fortran *contiguous* in memory, C-like order
+        order if ``a`` is Fortran *contiguous* in memory, C-like order
         otherwise.
+    copy : bool, optional
+        If ``True``, then the array data is copied. If ``None``, a copy will
+        only be made if it's required by ``order``. For ``False`` it raises
+        a ``ValueError`` if a copy cannot be avoided. Default: ``None``.
 
     Returns
     -------
@@ -255,9 +262,9 @@ def reshape(a, newshape, order='C'):
     It is not always possible to change the shape of an array without copying
     the data.
 
-    The `order` keyword gives the index ordering both for *fetching* the values
-    from `a`, and then *placing* the values into the output array.
-    For example, let's say you have an array:
+    The ``order`` keyword gives the index ordering both for *fetching*
+    the values from ``a``, and then *placing* the values into the output
+    array. For example, let's say you have an array:
 
     >>> a = np.arange(6).reshape((3, 2))
     >>> a
@@ -296,7 +303,16 @@ def reshape(a, newshape, order='C'):
            [3, 4],
            [5, 6]])
     """
-    return _wrapfunc(a, 'reshape', newshape, order=order)
+    if newshape is None and shape is None:
+        raise TypeError(
+            "reshape() missing 1 required positional argument: 'shape'")
+    if newshape is not None and shape is not None:
+        raise ValueError(
+            "You cannot specify 'newshape' and 'shape' arguments "
+            "at the same time.")
+    if shape is not None:
+        newshape = shape
+    return _wrapfunc(a, 'reshape', newshape, order=order, copy=copy)
 
 
 def _choose_dispatcher(a, choices, out=None, mode=None):
diff --git a/numpy/_core/fromnumeric.pyi b/numpy/_core/fromnumeric.pyi
@@ -94,12 +94,14 @@ def reshape(
     a: _ArrayLike[_SCT],
     newshape: _ShapeLike,
     order: _OrderACF = ...,
+    copy: None | bool = ...,
 ) -> NDArray[_SCT]: ...
 @overload
 def reshape(
     a: ArrayLike,
     newshape: _ShapeLike,
     order: _OrderACF = ...,
+    copy: None | bool = ...,
 ) -> NDArray[Any]: ...
 
 @overload
diff --git a/numpy/_core/src/multiarray/methods.c b/numpy/_core/src/multiarray/methods.c
@@ -181,14 +181,16 @@ array_put(PyArrayObject *self, PyObject *args, PyObject *kwds)
 static PyObject *
 array_reshape(PyArrayObject *self, PyObject *args, PyObject *kwds)
 {
-    static char *keywords[] = {"order", NULL};
+    static char *keywords[] = {"order", "copy", NULL};
     PyArray_Dims newshape;
     PyObject *ret;
     NPY_ORDER order = NPY_CORDER;
+    NPY_COPYMODE copy = NPY_COPY_IF_NEEDED;
     Py_ssize_t n = PyTuple_Size(args);
 
-    if (!NpyArg_ParseKeywords(kwds, "|O&", keywords,
-                PyArray_OrderConverter, &order)) {
+    if (!NpyArg_ParseKeywords(kwds, "|$O&O&", keywords,
+                PyArray_OrderConverter, &order,
+                PyArray_CopyConverter, &copy)) {
         return NULL;
     }
 
@@ -210,7 +212,7 @@ array_reshape(PyArrayObject *self, PyObject *args, PyObject *kwds)
             goto fail;
         }
     }
-    ret = PyArray_Newshape(self, &newshape, order);
+    ret = _reshape_with_copy_arg(self, &newshape, order, copy);
     npy_free_cache_dim_obj(newshape);
     return ret;
 
diff --git a/numpy/_core/src/multiarray/shape.c b/numpy/_core/src/multiarray/shape.c
@@ -201,6 +201,14 @@ PyArray_Resize(PyArrayObject *self, PyArray_Dims *newshape, int refcheck,
 NPY_NO_EXPORT PyObject *
 PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims,
                  NPY_ORDER order)
+{
+    return _reshape_with_copy_arg(self, newdims, order, NPY_COPY_IF_NEEDED);
+}
+
+
+NPY_NO_EXPORT PyObject *
+_reshape_with_copy_arg(PyArrayObject *array, PyArray_Dims *newdims,
+                       NPY_ORDER order, NPY_COPYMODE copy)
 {
     npy_intp i;
     npy_intp *dimensions = newdims->ptr;
@@ -212,64 +220,78 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims,
     int flags;
 
     if (order == NPY_ANYORDER) {
-        order = PyArray_ISFORTRAN(self) ? NPY_FORTRANORDER : NPY_CORDER;
+        order = PyArray_ISFORTRAN(array) ? NPY_FORTRANORDER : NPY_CORDER;
     }
     else if (order == NPY_KEEPORDER) {
         PyErr_SetString(PyExc_ValueError,
                 "order 'K' is not permitted for reshaping");
         return NULL;
     }
     /*  Quick check to make sure anything actually needs to be done */
-    if (ndim == PyArray_NDIM(self)) {
+    if (ndim == PyArray_NDIM(array) && copy != NPY_COPY_ALWAYS) {
         same = NPY_TRUE;
         i = 0;
         while (same && i < ndim) {
-            if (PyArray_DIM(self,i) != dimensions[i]) {
+            if (PyArray_DIM(array, i) != dimensions[i]) {
                 same=NPY_FALSE;
             }
             i++;
         }
         if (same) {
-            return PyArray_View(self, NULL, NULL);
+            return PyArray_View(array, NULL, NULL);
         }
     }
 
     /*
      * fix any -1 dimensions and check new-dimensions against old size
      */
-    if (_fix_unknown_dimension(newdims, self) < 0) {
+    if (_fix_unknown_dimension(newdims, array) < 0) {
         return NULL;
     }
-    /*
-     * sometimes we have to create a new copy of the array
-     * in order to get the right orientation and
-     * because we can't just reuse the buffer with the
-     * data in the order it is in.
-     */
-    Py_INCREF(self);
-    if (((order == NPY_CORDER && !PyArray_IS_C_CONTIGUOUS(self)) ||
-         (order == NPY_FORTRANORDER && !PyArray_IS_F_CONTIGUOUS(self)))) {
-        int success = 0;
-        success = _attempt_nocopy_reshape(self, ndim, dimensions,
-                                          newstrides, order);
-        if (success) {
-            /* no need to copy the array after all */
-            strides = newstrides;
+    if (copy == NPY_COPY_ALWAYS) {
+        PyObject *newcopy = PyArray_NewCopy(array, order);
+        if (newcopy == NULL) {
+            return NULL;
         }
-        else {
-            PyObject *newcopy;
-            newcopy = PyArray_NewCopy(self, order);
-            Py_DECREF(self);
-            if (newcopy == NULL) {
+        array = (PyArrayObject *)newcopy;
+    }
+    else {
+        /*
+         * sometimes we have to create a new copy of the array
+         * in order to get the right orientation and
+         * because we can't just reuse the buffer with the
+         * data in the order it is in.
+         */
+        Py_INCREF(array);
+        if (((order == NPY_CORDER && !PyArray_IS_C_CONTIGUOUS(array)) ||
+                (order == NPY_FORTRANORDER && !PyArray_IS_F_CONTIGUOUS(array)))) {
+            int success = 0;
+            success = _attempt_nocopy_reshape(array, ndim, dimensions,
+                                              newstrides, order);
+            if (success) {
+                /* no need to copy the array after all */
+                strides = newstrides;
+            }
+            else if (copy == NPY_COPY_NEVER) {
+                PyErr_SetString(PyExc_ValueError,
+                                "Unable to avoid creating a copy while reshaping.");
+                Py_DECREF(array);
                 return NULL;
             }
-            self = (PyArrayObject *)newcopy;
+            else {
+                PyObject *newcopy = PyArray_NewCopy(array, order);
+                Py_DECREF(array);
+                if (newcopy == NULL) {
+                    return NULL;
+                }
+                array = (PyArrayObject *)newcopy;
+            }
         }
     }
     /* We always have to interpret the contiguous buffer correctly */
 
     /* Make sure the flags argument is set. */
-    flags = PyArray_FLAGS(self);
+    flags = PyArray_FLAGS(array);
     if (ndim > 1) {
         if (order == NPY_FORTRANORDER) {
             flags &= ~NPY_ARRAY_C_CONTIGUOUS;
@@ -281,18 +303,17 @@ PyArray_Newshape(PyArrayObject *self, PyArray_Dims *newdims,
         }
     }
 
-    Py_INCREF(PyArray_DESCR(self));
+    Py_INCREF(PyArray_DESCR(array));
     ret = (PyArrayObject *)PyArray_NewFromDescr_int(
-            Py_TYPE(self), PyArray_DESCR(self),
-            ndim, dimensions, strides, PyArray_DATA(self),
-            flags, (PyObject *)self, (PyObject *)self,
+            Py_TYPE(array), PyArray_DESCR(array),
+            ndim, dimensions, strides, PyArray_DATA(array),
+            flags, (PyObject *)array, (PyObject *)array,
             _NPY_ARRAY_ENSURE_DTYPE_IDENTITY);
-    Py_DECREF(self);
+    Py_DECREF(array);
     return (PyObject *)ret;
 }
 
 
-
 /* For backward compatibility -- Not recommended */
 
 /*NUMPY_API
diff --git a/numpy/_core/src/multiarray/shape.h b/numpy/_core/src/multiarray/shape.h
@@ -1,6 +1,8 @@
 #ifndef NUMPY_CORE_SRC_MULTIARRAY_SHAPE_H_
 #define NUMPY_CORE_SRC_MULTIARRAY_SHAPE_H_
 
+#include "conversion_utils.h"
+
 /*
  * Creates a sorted stride perm matching the KEEPORDER behavior
  * of the NpyIter object. Because this operates based on multiple
@@ -27,4 +29,8 @@ PyArray_SqueezeSelected(PyArrayObject *self, npy_bool *axis_flags);
 NPY_NO_EXPORT PyObject *
 PyArray_MatrixTranspose(PyArrayObject *ap);
 
+NPY_NO_EXPORT PyObject *
+_reshape_with_copy_arg(PyArrayObject *array, PyArray_Dims *newdims,
+                       NPY_ORDER order, NPY_COPYMODE copy);
+
 #endif  /* NUMPY_CORE_SRC_MULTIARRAY_SHAPE_H_ */
diff --git a/numpy/_core/tests/test_numeric.py b/numpy/_core/tests/test_numeric.py
@@ -165,6 +165,52 @@ def test_reshape(self):
         tgt = [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]]
         assert_equal(np.reshape(arr, (2, 6)), tgt)
 
+    def test_reshape_shape_arg(self):
+        arr = np.arange(12)
+        shape = (3, 4)
+        expected = arr.reshape(shape)
+
+        with pytest.raises(
+            ValueError,
+            match="You cannot specify 'newshape' and 'shape' "
+                  "arguments at the same time."
+        ):
+            np.reshape(arr, newshape=shape, shape=shape)
+        with pytest.raises(
+            TypeError,
+            match=r"reshape\(\) missing 1 required positional "
+                  "argument: 'shape'"
+        ):
+            np.reshape(arr)
+
+        assert_equal(np.reshape(arr, shape), expected)
+        assert_equal(np.reshape(arr, shape, order="C"), expected)
+        assert_equal(np.reshape(arr, newshape=shape), expected)
+        assert_equal(np.reshape(arr, shape=shape), expected)
+        assert_equal(np.reshape(arr, shape=shape, order="C"), expected)
+
+    def test_reshape_copy_arg(self):
+        arr = np.arange(24).reshape(2, 3, 4)
+        arr_f_ord = np.array(arr, order="F")
+        shape = (12, 2)
+
+        assert np.shares_memory(np.reshape(arr, shape), arr)
+        assert np.shares_memory(np.reshape(arr, shape, order="C"), arr)
+        assert np.shares_memory(np.reshape(arr_f_ord, shape, order="F"), arr_f_ord)
+        assert np.shares_memory(np.reshape(arr, shape, copy=None), arr)
+        assert np.shares_memory(np.reshape(arr, shape, copy=False), arr)
+        assert np.shares_memory(arr.reshape(shape, copy=False), arr)
+        assert not np.shares_memory(np.reshape(arr, shape, copy=True), arr)
+        assert not np.shares_memory(np.reshape(arr, shape, order="C", copy=True), arr)
+        assert not np.shares_memory(np.reshape(arr, shape, order="F", copy=True), arr)
+        assert not np.shares_memory(np.reshape(arr, shape, order="F", copy=None), arr)
+
+        err_msg = "Unable to avoid creating a copy while reshaping."
+        with pytest.raises(ValueError, match=err_msg):
+            np.reshape(arr, shape, order="F", copy=False)
+        with pytest.raises(ValueError, match=err_msg):
+            np.reshape(arr_f_ord, shape, order="C", copy=False)
+
     def test_round(self):
         arr = [1.56, 72.54, 6.35, 3.25]
         tgt = [1.6, 72.5, 6.4, 3.2]
diff --git a/tools/ci/array-api-skips.txt b/tools/ci/array-api-skips.txt
@@ -7,7 +7,7 @@ array_api_tests/test_operators_and_elementwise_functions.py::test_ceil
 array_api_tests/test_operators_and_elementwise_functions.py::test_floor
 array_api_tests/test_operators_and_elementwise_functions.py::test_trunc
 
-# 'newshape' should be named 'shape'
+# 'shape' arg is present. 'newshape' is retained for backward compat.
 array_api_tests/test_signatures.py::test_func_signature[reshape]
 
 # missing 'descending' keyword arguments
