diff options
Diffstat (limited to 'numpy')
-rw-r--r-- | numpy/core/src/multiarray/item_selection.c | 47 | ||||
-rw-r--r-- | numpy/core/tests/test_numeric.py | 13 | ||||
-rw-r--r-- | numpy/core/tests/test_regression.py | 5 | ||||
-rw-r--r-- | numpy/matrixlib/tests/test_defmatrix.py | 4 | ||||
-rw-r--r-- | numpy/matrixlib/tests/test_regression.py | 4 |
5 files changed, 39 insertions, 34 deletions
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c index 00639a19c..9aaea1bfb 100644 --- a/numpy/core/src/multiarray/item_selection.c +++ b/numpy/core/src/multiarray/item_selection.c @@ -2331,36 +2331,27 @@ finish: } /* Create views into ret, one for each dimension */ - if (ndim == 1) { - /* Directly switch to one dimensions (dimension 1 is 1 anyway) */ - ((PyArrayObject_fields *)ret)->nd = 1; - PyTuple_SET_ITEM(ret_tuple, 0, (PyObject *)ret); - } - else { - for (i = 0; i < ndim; ++i) { - PyArrayObject *view; - npy_intp stride = ndim * NPY_SIZEOF_INTP; - - view = (PyArrayObject *)PyArray_New(Py_TYPE(self), 1, - &nonzero_count, - NPY_INTP, &stride, - PyArray_BYTES(ret) + i*NPY_SIZEOF_INTP, - 0, 0, (PyObject *)self); - if (view == NULL) { - Py_DECREF(ret); - Py_DECREF(ret_tuple); - return NULL; - } - Py_INCREF(ret); - if (PyArray_SetBaseObject(view, (PyObject *)ret) < 0) { - Py_DECREF(ret); - Py_DECREF(ret_tuple); - } - PyTuple_SET_ITEM(ret_tuple, i, (PyObject *)view); + for (i = 0; i < ndim; ++i) { + npy_intp stride = ndim * NPY_SIZEOF_INTP; + + PyArrayObject *view = (PyArrayObject *)PyArray_New(Py_TYPE(ret), 1, + &nonzero_count, NPY_INTP, &stride, + PyArray_BYTES(ret) + i*NPY_SIZEOF_INTP, + 0, PyArray_FLAGS(ret), (PyObject *)ret); + if (view == NULL) { + Py_DECREF(ret); + Py_DECREF(ret_tuple); + return NULL; } - - Py_DECREF(ret); + Py_INCREF(ret); + if (PyArray_SetBaseObject(view, (PyObject *)ret) < 0) { + Py_DECREF(ret); + Py_DECREF(ret_tuple); + return NULL; + } + PyTuple_SET_ITEM(ret_tuple, i, (PyObject *)view); } + Py_DECREF(ret); return ret_tuple; } diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py index 9f0fb47e5..87b93ff3c 100644 --- a/numpy/core/tests/test_numeric.py +++ b/numpy/core/tests/test_numeric.py @@ -951,6 +951,19 @@ class TestNonzero(TestCase): assert_equal(np.nonzero(c)[0], np.concatenate((np.arange(10 +i, 20 + i), [20 +i*2]))) + def test_return_type(self): + class C(np.ndarray): + pass + + for view in (C, np.ndarray): + for nd in range(1, 4): + shape = tuple(range(2, 2+nd)) + x = np.arange(np.prod(shape)).reshape(shape).view(view) + for nzx in (np.nonzero(x), x.nonzero()): + for nzx_i in nzx: + assert_(type(nzx_i) is np.ndarray) + assert_(nzx_i.flags.writeable) + class TestIndex(TestCase): def test_boolean(self): diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py index 9e8511a01..76cfa8e0b 100644 --- a/numpy/core/tests/test_regression.py +++ b/numpy/core/tests/test_regression.py @@ -1122,8 +1122,6 @@ class TestRegression(TestCase): assert_(dat.mean(1).info == 'jubba') assert_(dat.min(1).info == 'jubba') assert_(dat.newbyteorder().info == 'jubba') - assert_(dat.nonzero()[0].info == 'jubba') - assert_(dat.nonzero()[1].info == 'jubba') assert_(dat.prod(1).info == 'jubba') assert_(dat.ptp(1).info == 'jubba') assert_(dat.ravel().info == 'jubba') @@ -1140,6 +1138,9 @@ class TestRegression(TestCase): assert_(dat.T.info == 'jubba') assert_(dat.var(1).info == 'jubba') assert_(dat.view(TestArray).info == 'jubba') + # These methods do not preserve subclasses + assert_(type(dat.nonzero()[0]) is np.ndarray) + assert_(type(dat.nonzero()[1]) is np.ndarray) def test_recarray_tolist(self, level=rlevel): """Ticket #793, changeset r5215 diff --git a/numpy/matrixlib/tests/test_defmatrix.py b/numpy/matrixlib/tests/test_defmatrix.py index f3a8e72ca..166f68862 100644 --- a/numpy/matrixlib/tests/test_defmatrix.py +++ b/numpy/matrixlib/tests/test_defmatrix.py @@ -306,8 +306,8 @@ class TestMatrixReturn(TestCase): assert_(type(a.real) is matrix) assert_(type(a.imag) is matrix) c, d = matrix([0.0]).nonzero() - assert_(type(c) is matrix) - assert_(type(d) is matrix) + assert_(type(c) is np.ndarray) + assert_(type(d) is np.ndarray) class TestIndexing(TestCase): diff --git a/numpy/matrixlib/tests/test_regression.py b/numpy/matrixlib/tests/test_regression.py index 119b21d8a..40062653f 100644 --- a/numpy/matrixlib/tests/test_regression.py +++ b/numpy/matrixlib/tests/test_regression.py @@ -17,8 +17,8 @@ class TestRegression(TestCase): assert_(type(a.real) is np.matrix) assert_(type(a.imag) is np.matrix) c, d = np.matrix([0.0]).nonzero() - assert_(type(c) is np.matrix) - assert_(type(d) is np.matrix) + assert_(type(c) is np.ndarray) + assert_(type(d) is np.ndarray) def test_matrix_multiply_by_1d_vector(self, level=rlevel) : """Ticket #473""" |