diff options
| -rw-r--r-- | numpy/core/src/multiarray/item_selection.c | 13 | ||||
| -rw-r--r-- | numpy/core/tests/test_multiarray.py | 3 |
2 files changed, 11 insertions, 5 deletions
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c index 80dc8201f..dcd3322c4 100644 --- a/numpy/core/src/multiarray/item_selection.c +++ b/numpy/core/src/multiarray/item_selection.c @@ -1441,11 +1441,6 @@ PyArray_LexSort(PyObject *sort_keys, int axis) && PyDataType_FLAGCHK(PyArray_DESCR(mps[i]), NPY_NEEDS_PYAPI)) { object = 1; } - its[i] = (PyArrayIterObject *)PyArray_IterAllButAxis( - (PyObject *)mps[i], &axis); - if (its[i] == NULL) { - goto fail; - } } /* Now we can check the axis */ @@ -1472,6 +1467,14 @@ PyArray_LexSort(PyObject *sort_keys, int axis) goto fail; } + for (i = 0; i < n; i++) { + its[i] = (PyArrayIterObject *)PyArray_IterAllButAxis( + (PyObject *)mps[i], &axis); + if (its[i] == NULL) { + goto fail; + } + } + /* Now do the sorting */ ret = (PyArrayObject *)PyArray_New(&PyArray_Type, PyArray_NDIM(mps[0]), PyArray_DIMS(mps[0]), NPY_INTP, diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py index 8cd28f88f..4a2a232af 100644 --- a/numpy/core/tests/test_multiarray.py +++ b/numpy/core/tests/test_multiarray.py @@ -3614,6 +3614,9 @@ class TestLexsort(TestCase): u, v = np.array(u, dtype='object'), np.array(v, dtype='object') assert_array_equal(idx, np.lexsort((u, v))) + def test_invalid_axis(self): # gh-7528 + x = np.linspace(0., 1., 42*3).reshape(42, 3) + assert_raises(ValueError, np.lexsort, x, axis=2) class TestIO(object): """Test tofile, fromfile, tobytes, and fromstring""" |