diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/function_base.py | 22 | ||||
| -rw-r--r-- | numpy/lib/tests/test_function_base.py | 18 |
2 files changed, 31 insertions, 9 deletions
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index b530f0aa1..556227c0d 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -1631,21 +1631,27 @@ def trim_zeros(filt, trim='fb'): # Numpy 1.20.0, 2020-07-31 warning = DeprecationWarning( "in the future trim_zeros will require a 1-D array as input " - "that is compatible with ndarray.astype(bool)" + "that supports elementwise comparisons with zero" ) warning.__cause__ = ex - warnings.warn(warning, stacklevel=3) - # Fall back to the old implementation if an exception is encountered - # Note that the same exception may or may not be raised here as well - return _trim_zeros_old(filt, trim) + # Fall back to the old implementation if an exception is encountered + # Note that the same exception may or may not be raised here as well + ret = _trim_zeros_old(filt, trim) + warnings.warn(warning, stacklevel=3) + return ret def _trim_zeros_new(filt, trim='fb'): """Newer optimized implementation of ``trim_zeros()``.""" - arr = np.asanyarray(filt).astype(bool, copy=False) - - if arr.ndim != 1: + arr_any = np.asanyarray(filt) + arr = arr_any != 0 if arr_any.dtype != bool else arr_any + + if arr is False: + # not all dtypes support elementwise comparisons with `0` (e.g. str); + # they will return `False` instead + raise TypeError('elementwise comparison failed; unsupported data type') + elif arr.ndim != 1: raise ValueError('trim_zeros requires an array of exactly one dimension') elif not len(arr): return filt diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 635fe1432..34a395ee4 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -1169,7 +1169,7 @@ class TestTrimZeros: a = np.array([0, 0, 1, 0, 2, 3, 4, 0]) b = a.astype(float) c = a.astype(complex) - d = np.array([None, [], 1, False, 'b', 3.0, range(4), b''], dtype=object) + d = a.astype(object) def values(self): attr_names = ('a', 'b', 'c', 'd') @@ -1208,6 +1208,22 @@ class TestTrimZeros: res = trim_zeros(arr) assert_array_equal(arr, res) + @pytest.mark.parametrize( + 'arr', + [np.array([0, 2**62, 0]), + np.array([0, 2**63, 0]), + np.array([0, 2**64, 0])] + ) + def test_overflow(self, arr): + slc = np.s_[1:2] + res = trim_zeros(arr) + assert_array_equal(res, arr[slc]) + + def test_no_trim(self): + arr = np.array([None, 1, None]) + res = trim_zeros(arr) + assert_array_equal(arr, res) + class TestExtins: |