diff options
Diffstat (limited to 'numpy/lib')
| -rw-r--r-- | numpy/lib/__init__.py | 2 | ||||
| -rw-r--r-- | numpy/lib/arraysetops.py | 19 | ||||
| -rw-r--r-- | numpy/lib/tests/test_arraysetops.py | 67 |
3 files changed, 84 insertions, 4 deletions
diff --git a/numpy/lib/__init__.py b/numpy/lib/__init__.py index dc40ac67b..c1757150e 100644 --- a/numpy/lib/__init__.py +++ b/numpy/lib/__init__.py @@ -26,7 +26,7 @@ from .financial import * from .arrayterator import Arrayterator from .arraypad import * from ._version import * -from numpy.core.multiarray import tracemalloc_domain +from numpy.core._multiarray_umath import tracemalloc_domain __all__ = ['emath', 'math', 'tracemalloc_domain'] __all__ += type_check.__all__ diff --git a/numpy/lib/arraysetops.py b/numpy/lib/arraysetops.py index d84455a8f..62e9b6d50 100644 --- a/numpy/lib/arraysetops.py +++ b/numpy/lib/arraysetops.py @@ -82,6 +82,11 @@ def ediff1d(ary, to_end=None, to_begin=None): # force a 1d array ary = np.asanyarray(ary).ravel() + # we have unit tests enforcing + # propagation of the dtype of input + # ary to returned result + dtype_req = ary.dtype + # fast track default case if to_begin is None and to_end is None: return ary[1:] - ary[:-1] @@ -89,13 +94,23 @@ def ediff1d(ary, to_end=None, to_begin=None): if to_begin is None: l_begin = 0 else: - to_begin = np.asanyarray(to_begin).ravel() + to_begin = np.asanyarray(to_begin) + if not np.can_cast(to_begin, dtype_req): + raise TypeError("dtype of to_begin must be compatible " + "with input ary") + + to_begin = to_begin.ravel() l_begin = len(to_begin) if to_end is None: l_end = 0 else: - to_end = np.asanyarray(to_end).ravel() + to_end = np.asanyarray(to_end) + if not np.can_cast(to_end, dtype_req): + raise TypeError("dtype of to_end must be compatible " + "with input ary") + + to_end = to_end.ravel() l_end = len(to_end) # do the calculation in place and copy to_begin and to_end diff --git a/numpy/lib/tests/test_arraysetops.py b/numpy/lib/tests/test_arraysetops.py index c76afb8e5..4b61726d2 100644 --- a/numpy/lib/tests/test_arraysetops.py +++ b/numpy/lib/tests/test_arraysetops.py @@ -6,10 +6,13 @@ from __future__ import division, absolute_import, print_function import numpy as np import sys -from numpy.testing import assert_array_equal, assert_equal, assert_raises +from numpy.testing import (assert_array_equal, assert_equal, + assert_raises, assert_raises_regex) from numpy.lib.arraysetops import ( ediff1d, intersect1d, setxor1d, union1d, setdiff1d, unique, in1d, isin ) +import pytest + class TestSetOps(object): @@ -125,6 +128,68 @@ class TestSetOps(object): assert_array_equal([7,1], ediff1d(two_elem, to_begin=7)) assert_array_equal([5,6,1], ediff1d(two_elem, to_begin=[5,6])) + @pytest.mark.parametrize("ary, prepend, append", [ + # should fail because trying to cast + # np.nan standard floating point value + # into an integer array: + (np.array([1, 2, 3], dtype=np.int64), + None, + np.nan), + # should fail because attempting + # to downcast to smaller int type: + (np.array([1, 2, 3], dtype=np.int32), + np.array([5, 7, 2], dtype=np.int64), + None), + # should fail because attempting to cast + # two special floating point values + # to integers (on both sides of ary): + (np.array([1., 3., 9.], dtype=np.int8), + np.nan, + np.nan), + ]) + def test_ediff1d_forbidden_type_casts(self, ary, prepend, append): + # verify resolution of gh-11490 + + # specifically, raise an appropriate + # Exception when attempting to append or + # prepend with an incompatible type + msg = 'must be compatible' + with assert_raises_regex(TypeError, msg): + ediff1d(ary=ary, + to_end=append, + to_begin=prepend) + + @pytest.mark.parametrize("ary," + "prepend," + "append," + "expected", [ + (np.array([1, 2, 3], dtype=np.int16), + 0, + None, + np.array([0, 1, 1], dtype=np.int16)), + (np.array([1, 2, 3], dtype=np.int32), + 0, + 0, + np.array([0, 1, 1, 0], dtype=np.int32)), + (np.array([1, 2, 3], dtype=np.int64), + 3, + -9, + np.array([3, 1, 1, -9], dtype=np.int64)), + ]) + def test_ediff1d_scalar_handling(self, + ary, + prepend, + append, + expected): + # maintain backwards-compatibility + # of scalar prepend / append behavior + # in ediff1d following fix for gh-11490 + actual = np.ediff1d(ary=ary, + to_end=append, + to_begin=prepend) + assert_equal(actual, expected) + + def test_isin(self): # the tests for in1d cover most of isin's behavior # if in1d is removed, would need to change those tests to test |