diff options
| author | Eric Wieser <wieser.eric@gmail.com> | 2018-01-31 09:38:14 -0800 |
|---|---|---|
| committer | Eric Wieser <wieser.eric@gmail.com> | 2018-02-03 11:36:30 -0800 |
| commit | ef70f13177a53266fd8547da6e00bc252a057893 (patch) | |
| tree | b836e830ac216985e27c3f6e63e582df2bbab91c /numpy/lib | |
| parent | 2854d508d1c6d211f2ce99e8747eda1cb427a78a (diff) | |
| download | numpy-ef70f13177a53266fd8547da6e00bc252a057893.tar.gz | |
MAINT: Use AxisError in swapaxes
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/arraysetops.py | 8 | ||||
| -rw-r--r-- | numpy/lib/tests/test_arraysetops.py | 4 |
2 files changed, 7 insertions, 5 deletions
diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index e6ff5bf38..7b103ef3e 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -223,10 +223,12 @@ def unique(ar, return_index=False, return_inverse=False, ret = _unique1d(ar, return_index, return_inverse, return_counts) return _unpack_tuple(ret) - if not (-ar.ndim <= axis < ar.ndim): - raise ValueError('Invalid axis kwarg specified for unique') + try: + ar = np.swapaxes(ar, axis, 0) + except np.AxisError: + # this removes the "axis1" or "axis2" prefix from the error message + raise np.AxisError(axis, ar.ndim) - ar = np.swapaxes(ar, axis, 0) orig_shape, orig_dtype = ar.shape, ar.dtype # Must reshape to a contiguous 2D array for this to work... ar = ar.reshape(orig_shape[0], -1) diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py index c2ba7ac86..17415d8fe 100644 --- a/numpy/lib/tests/test_arraysetops.py +++ b/numpy/lib/tests/test_arraysetops.py @@ -409,8 +409,8 @@ class TestUnique(object): assert_raises(TypeError, self._run_axis_tests, [('a', int), ('b', object)]) - assert_raises(ValueError, unique, np.arange(10), axis=2) - assert_raises(ValueError, unique, np.arange(10), axis=-2) + assert_raises(np.AxisError, unique, np.arange(10), axis=2) + assert_raises(np.AxisError, unique, np.arange(10), axis=-2) def test_unique_axis_list(self): msg = "Unique failed on list of lists" |