import sys
|
|
import pytest
|
|
import numpy as np
|
from numpy.testing import HAS_REFCOUNT, assert_, assert_array_equal, assert_raises
|
|
|
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((), 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)]
|
if HAS_REFCOUNT:
|
orig_rcs = [sys.getrefcount(o) for o in objects]
|
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) == rc + 1
|
for o, rc in zip(objects, orig_rcs)))
|
# 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) == rc + 1
|
for o, rc in zip(objects, orig_rcs)))
|
|
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)
|