Merge pull request #7442 from jaimefrio/refactor_diagonal

MANT: Simplify diagonal length calculation logic

1 files changed, 12 insertions, 26 deletions

diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
index 9789235c2..80dc8201f 100644
--- a/numpy/core/src/multiarray/item_selection.c
+++ b/numpy/core/src/multiarray/item_selection.c
@@ -1813,7 +1813,7 @@ PyArray_Diagonal(PyArrayObject *self, int offset, int axis1, int axis2)
 {
     int i, idim, ndim = PyArray_NDIM(self);
     npy_intp *strides;
-    npy_intp stride1, stride2;
+    npy_intp stride1, stride2, offset_stride;
     npy_intp *shape, dim1, dim2;
 
     char *data;
@@ -1860,35 +1860,21 @@ PyArray_Diagonal(PyArrayObject *self, int offset, int axis1, int axis2)
 
     /* Compute the data pointers and diag_size for the view */
     data = PyArray_DATA(self);
-    if (offset > 0) {
-        if (offset >= dim2) {
-            diag_size = 0;
-        }
-        else {
-            data += offset * stride2;
-
-            diag_size = dim2 - offset;
-            if (dim1 < diag_size) {
-                diag_size = dim1;
-            }
-        }
+    if (offset >= 0) {
+        offset_stride = stride2;
+        dim2 -= offset;
     }
-    else if (offset < 0) {
+    else {
         offset = -offset;
-        if (offset >= dim1) {
-            diag_size = 0;
-        }
-        else {
-            data += offset * stride1;
-
-            diag_size = dim1 - offset;
-            if (dim2 < diag_size) {
-                diag_size = dim2;
-            }
-        }
+        offset_stride = stride1;
+        dim1 -= offset;
+    }
+    diag_size = dim2 < dim1 ? dim2 : dim1;
+    if (diag_size < 0) {
+        diag_size = 0;
     }
     else {
-        diag_size = dim1 < dim2 ? dim1 : dim2;
+        data += offset * offset_stride;
     }
 
     /* Build the new shape and strides for the main data */