name : test_item_selection.py import sys import pytest import numpy as np from numpy.testing import ( assert_, assert_raises, assert_array_equal, HAS_REFCOUNT ) class TestTake: def test_simple(self): a = [[1, 2], [3, 4]] a_str = [[b'1', b'2'], [b'3', b'4']] modes = ['raise', 'wrap', 'clip'] indices = [-1, 4] index_arrays = [np.empty(0, dtype=np.intp), np.empty(tuple(), dtype=np.intp), np.empty((1, 1), dtype=np.intp)] real_indices = {'raise': {-1: 1, 4: IndexError}, 'wrap': {-1: 1, 4: 0}, 'clip': {-1: 0, 4: 1}} # Currently all types but object, use the same function generation. # So it should not be necessary to test all. However test also a non # refcounted struct on top of object, which has a size that hits the # default (non-specialized) path. types = int, object, np.dtype([('', 'i2', 3)]) for t in types: # ta works, even if the array may be odd if buffer interface is used ta = np.array(a if np.issubdtype(t, np.number) else a_str, dtype=t) tresult = list(ta.T.copy()) for index_array in index_arrays: if index_array.size != 0: tresult[0].shape = (2,) + index_array.shape tresult[1].shape = (2,) + index_array.shape for mode in modes: for index in indices: real_index = real_indices[mode][index] if real_index is IndexError and index_array.size != 0: index_array.put(0, index) assert_raises(IndexError, ta.take, index_array, mode=mode, axis=1) elif index_array.size != 0: index_array.put(0, index) res = ta.take(index_array, mode=mode, axis=1) assert_array_equal(res, tresult[real_index]) else: res = ta.take(index_array, mode=mode, axis=1) assert_(res.shape == (2,) + index_array.shape) def test_refcounting(self): objects = [object() for i in range(10)] for mode in ('raise', 'clip', 'wrap'): a = np.array(objects) b = np.array([2, 2, 4, 5, 3, 5]) a.take(b, out=a[:6], mode=mode) del a if HAS_REFCOUNT: assert_(all(sys.getrefcount(o) == 3 for o in objects)) # not contiguous, example: a = np.array(objects * 2)[::2] a.take(b, out=a[:6], mode=mode) del a if HAS_REFCOUNT: assert_(all(sys.getrefcount(o) == 3 for o in objects)) def test_unicode_mode(self): d = np.arange(10) k = b'\xc3\xa4'.decode("UTF8") assert_raises(ValueError, d.take, 5, mode=k) def test_empty_partition(self): # In reference to github issue #6530 a_original = np.array([0, 2, 4, 6, 8, 10]) a = a_original.copy() # An empty partition should be a successful no-op a.partition(np.array([], dtype=np.int16)) assert_array_equal(a, a_original) def test_empty_argpartition(self): # In reference to github issue #6530 a = np.array([0, 2, 4, 6, 8, 10]) a = a.argpartition(np.array([], dtype=np.int16)) b = np.array([0, 1, 2, 3, 4, 5]) assert_array_equal(a, b) class TestPutMask: @pytest.mark.parametrize("dtype", list(np.typecodes["All"]) + ["i,O"]) def test_simple(self, dtype): if dtype.lower() == "m": dtype += "8[ns]" # putmask is weird and doesn't care about value length (even shorter) vals = np.arange(1001).astype(dtype=dtype) mask = np.random.randint(2, size=1000).astype(bool) # Use vals.dtype in case of flexible dtype (i.e. string) arr = np.zeros(1000, dtype=vals.dtype) zeros = arr.copy() np.putmask(arr, mask, vals) assert_array_equal(arr[mask], vals[:len(mask)][mask]) assert_array_equal(arr[~mask], zeros[~mask]) @pytest.mark.parametrize("dtype", list(np.typecodes["All"])[1:] + ["i,O"]) @pytest.mark.parametrize("mode", ["raise", "wrap", "clip"]) def test_empty(self, dtype, mode): arr = np.zeros(1000, dtype=dtype) arr_copy = arr.copy() mask = np.random.randint(2, size=1000).astype(bool) # Allowing empty values like this is weird... np.put(arr, mask, []) assert_array_equal(arr, arr_copy) class TestPut: @pytest.mark.parametrize("dtype", list(np.typecodes["All"])[1:] + ["i,O"]) @pytest.mark.parametrize("mode", ["raise", "wrap", "clip"]) def test_simple(self, dtype, mode): if dtype.lower() == "m": dtype += "8[ns]" # put is weird and doesn't care about value length (even shorter) vals = np.arange(1001).astype(dtype=dtype) # Use vals.dtype in case of flexible dtype (i.e. string) arr = np.zeros(1000, dtype=vals.dtype) zeros = arr.copy() if mode == "clip": # Special because 0 and -1 value are "reserved" for clip test indx = np.random.permutation(len(arr) - 2)[:-500] + 1 indx[-1] = 0 indx[-2] = len(arr) - 1 indx_put = indx.copy() indx_put[-1] = -1389 indx_put[-2] = 1321 else: # Avoid duplicates (for simplicity) and fill half only indx = np.random.permutation(len(arr) - 3)[:-500] indx_put = indx if mode == "wrap": indx_put = indx_put + len(arr) np.put(arr, indx_put, vals, mode=mode) assert_array_equal(arr[indx], vals[:len(indx)]) untouched = np.ones(len(arr), dtype=bool) untouched[indx] = False assert_array_equal(arr[untouched], zeros[:untouched.sum()]) @pytest.mark.parametrize("dtype", list(np.typecodes["All"])[1:] + ["i,O"]) @pytest.mark.parametrize("mode", ["raise", "wrap", "clip"]) def test_empty(self, dtype, mode): arr = np.zeros(1000, dtype=dtype) arr_copy = arr.copy() # Allowing empty values like this is weird... np.put(arr, [1, 2, 3], []) assert_array_equal(arr, arr_copy)