BUG: Guarantee array is in valid state after memory error occurs in getset.c (#15736)

Currently, in function array_shape_set, ndarray's pointers to dimensions and strides are freed before new array is allocated (Line 71). In case memory error occur, the array is left with dangling pointers. Therefore, we can not recover from such error.

* To guarantee array in valid state when memory error occur
* Fix: Free cache when nd == 0 to avoid memory leak
* Update numpy/core/src/multiarray/getset.c
Co-Authored-By: Eric Wieser <wieser.eric@gmail.com>

1 files changed, 13 insertions, 7 deletions

diff --git a/numpy/core/src/multiarray/getset.c b/numpy/core/src/multiarray/getset.c
index 5b7362153..8c1b7f943 100644
--- a/numpy/core/src/multiarray/getset.c
+++ b/numpy/core/src/multiarray/getset.c
@@ -67,28 +67,34 @@ array_shape_set(PyArrayObject *self, PyObject *val)
         return -1;
     }
 
-    /* Free old dimensions and strides */
-    npy_free_cache_dim_array(self);
     nd = PyArray_NDIM(ret);
-    ((PyArrayObject_fields *)self)->nd = nd;
     if (nd > 0) {
         /* create new dimensions and strides */
-        ((PyArrayObject_fields *)self)->dimensions = npy_alloc_cache_dim(2 * nd);
-        if (PyArray_DIMS(self) == NULL) {
+        npy_intp *_dimensions = npy_alloc_cache_dim(2 * nd);
+        if (_dimensions == NULL) {
             Py_DECREF(ret);
-            PyErr_SetString(PyExc_MemoryError,"");
+            PyErr_NoMemory();
             return -1;
         }
-        ((PyArrayObject_fields *)self)->strides = PyArray_DIMS(self) + nd;
+        /* Free old dimensions and strides */
+        npy_free_cache_dim_array(self);
+        ((PyArrayObject_fields *)self)->nd = nd;
+        ((PyArrayObject_fields *)self)->dimensions = _dimensions; 
+        ((PyArrayObject_fields *)self)->strides = _dimensions + nd;
+
         if (nd) {
             memcpy(PyArray_DIMS(self), PyArray_DIMS(ret), nd*sizeof(npy_intp));
             memcpy(PyArray_STRIDES(self), PyArray_STRIDES(ret), nd*sizeof(npy_intp));
         }
     }
     else {
+        /* Free old dimensions and strides */
+        npy_free_cache_dim_array(self);        
+        ((PyArrayObject_fields *)self)->nd = 0;
         ((PyArrayObject_fields *)self)->dimensions = NULL;
         ((PyArrayObject_fields *)self)->strides = NULL;
     }
+
     Py_DECREF(ret);
     PyArray_UpdateFlags(self, NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_F_CONTIGUOUS);
     return 0;