From 13111ad783b195aa66d262b679258a6b94fa1e07 Mon Sep 17 00:00:00 2001 From: gfyoung Date: Fri, 22 Jun 2018 00:45:26 -0700 Subject: [PATCH] TST: Clean up tests in test_take.py Utilizes pytest's fixtures to cleanup the code significantly. --- pandas/tests/test_take.py | 525 +++++++++++++++++--------------------- 1 file changed, 239 insertions(+), 286 deletions(-) diff --git a/pandas/tests/test_take.py b/pandas/tests/test_take.py index 9ab147edb8d1b..ade847923c083 100644 --- a/pandas/tests/test_take.py +++ b/pandas/tests/test_take.py @@ -10,315 +10,268 @@ from pandas._libs.tslib import iNaT +@pytest.fixture(params=[True, False]) +def writeable(request): + return request.param + + +# Check that take_nd works both with writeable arrays +# (in which case fast typed memory-views implementation) +# and read-only arrays alike. +@pytest.fixture(params=[ + (np.float64, True), + (np.float32, True), + (np.uint64, False), + (np.uint32, False), + (np.uint16, False), + (np.uint8, False), + (np.int64, False), + (np.int32, False), + (np.int16, False), + (np.int8, False), + (np.object_, True), + (np.bool, False), +]) +def dtype_can_hold_na(request): + return request.param + + +@pytest.fixture(params=[ + (np.int8, np.int16(127), np.int8), + (np.int8, np.int16(128), np.int16), + (np.int32, 1, np.int32), + (np.int32, 2.0, np.float64), + (np.int32, 3.0 + 4.0j, np.complex128), + (np.int32, True, np.object_), + (np.int32, "", np.object_), + (np.float64, 1, np.float64), + (np.float64, 2.0, np.float64), + (np.float64, 3.0 + 4.0j, np.complex128), + (np.float64, True, np.object_), + (np.float64, "", np.object_), + (np.complex128, 1, np.complex128), + (np.complex128, 2.0, np.complex128), + (np.complex128, 3.0 + 4.0j, np.complex128), + (np.complex128, True, np.object_), + (np.complex128, "", np.object_), + (np.bool_, 1, np.object_), + (np.bool_, 2.0, np.object_), + (np.bool_, 3.0 + 4.0j, np.object_), + (np.bool_, True, np.bool_), + (np.bool_, '', np.object_), +]) +def dtype_fill_out_dtype(request): + return request.param + + class TestTake(object): - # standard incompatible fill error + # Standard incompatible fill error. fill_error = re.compile("Incompatible type for fill_value") - def test_1d_with_out(self): - def _test_dtype(dtype, can_hold_na, writeable=True): - data = np.random.randint(0, 2, 4).astype(dtype) - data.flags.writeable = writeable + def test_1d_with_out(self, dtype_can_hold_na, writeable): + dtype, can_hold_na = dtype_can_hold_na + + data = np.random.randint(0, 2, 4).astype(dtype) + data.flags.writeable = writeable + + indexer = [2, 1, 0, 1] + out = np.empty(4, dtype=dtype) + algos.take_1d(data, indexer, out=out) - indexer = [2, 1, 0, 1] - out = np.empty(4, dtype=dtype) + expected = data.take(indexer) + tm.assert_almost_equal(out, expected) + + indexer = [2, 1, 0, -1] + out = np.empty(4, dtype=dtype) + + if can_hold_na: algos.take_1d(data, indexer, out=out) expected = data.take(indexer) + expected[3] = np.nan tm.assert_almost_equal(out, expected) - - indexer = [2, 1, 0, -1] - out = np.empty(4, dtype=dtype) - if can_hold_na: + else: + with tm.assert_raises_regex(TypeError, self.fill_error): algos.take_1d(data, indexer, out=out) - expected = data.take(indexer) - expected[3] = np.nan - tm.assert_almost_equal(out, expected) - else: - with tm.assert_raises_regex(TypeError, self.fill_error): - algos.take_1d(data, indexer, out=out) - # no exception o/w - data.take(indexer, out=out) - - for writeable in [True, False]: - # Check that take_nd works both with writeable arrays (in which - # case fast typed memoryviews implementation) and read-only - # arrays alike. - _test_dtype(np.float64, True, writeable=writeable) - _test_dtype(np.float32, True, writeable=writeable) - _test_dtype(np.uint64, False, writeable=writeable) - _test_dtype(np.uint32, False, writeable=writeable) - _test_dtype(np.uint16, False, writeable=writeable) - _test_dtype(np.uint8, False, writeable=writeable) - _test_dtype(np.int64, False, writeable=writeable) - _test_dtype(np.int32, False, writeable=writeable) - _test_dtype(np.int16, False, writeable=writeable) - _test_dtype(np.int8, False, writeable=writeable) - _test_dtype(np.object_, True, writeable=writeable) - _test_dtype(np.bool, False, writeable=writeable) - - def test_1d_fill_nonna(self): - def _test_dtype(dtype, fill_value, out_dtype): - data = np.random.randint(0, 2, 4).astype(dtype) - - indexer = [2, 1, 0, -1] - - result = algos.take_1d(data, indexer, fill_value=fill_value) - assert ((result[[0, 1, 2]] == data[[2, 1, 0]]).all()) - assert (result[3] == fill_value) - assert (result.dtype == out_dtype) - - indexer = [2, 1, 0, 1] - - result = algos.take_1d(data, indexer, fill_value=fill_value) - assert ((result[[0, 1, 2, 3]] == data[indexer]).all()) - assert (result.dtype == dtype) - - _test_dtype(np.int8, np.int16(127), np.int8) - _test_dtype(np.int8, np.int16(128), np.int16) - _test_dtype(np.int32, 1, np.int32) - _test_dtype(np.int32, 2.0, np.float64) - _test_dtype(np.int32, 3.0 + 4.0j, np.complex128) - _test_dtype(np.int32, True, np.object_) - _test_dtype(np.int32, '', np.object_) - _test_dtype(np.float64, 1, np.float64) - _test_dtype(np.float64, 2.0, np.float64) - _test_dtype(np.float64, 3.0 + 4.0j, np.complex128) - _test_dtype(np.float64, True, np.object_) - _test_dtype(np.float64, '', np.object_) - _test_dtype(np.complex128, 1, np.complex128) - _test_dtype(np.complex128, 2.0, np.complex128) - _test_dtype(np.complex128, 3.0 + 4.0j, np.complex128) - _test_dtype(np.complex128, True, np.object_) - _test_dtype(np.complex128, '', np.object_) - _test_dtype(np.bool_, 1, np.object_) - _test_dtype(np.bool_, 2.0, np.object_) - _test_dtype(np.bool_, 3.0 + 4.0j, np.object_) - _test_dtype(np.bool_, True, np.bool_) - _test_dtype(np.bool_, '', np.object_) - - def test_2d_with_out(self): - def _test_dtype(dtype, can_hold_na, writeable=True): - data = np.random.randint(0, 2, (5, 3)).astype(dtype) - data.flags.writeable = writeable - - indexer = [2, 1, 0, 1] - out0 = np.empty((4, 3), dtype=dtype) - out1 = np.empty((5, 4), dtype=dtype) + + # No Exception otherwise. + data.take(indexer, out=out) + + def test_1d_fill_nonna(self, dtype_fill_out_dtype): + dtype, fill_value, out_dtype = dtype_fill_out_dtype + data = np.random.randint(0, 2, 4).astype(dtype) + indexer = [2, 1, 0, -1] + + result = algos.take_1d(data, indexer, fill_value=fill_value) + assert ((result[[0, 1, 2]] == data[[2, 1, 0]]).all()) + assert (result[3] == fill_value) + assert (result.dtype == out_dtype) + + indexer = [2, 1, 0, 1] + + result = algos.take_1d(data, indexer, fill_value=fill_value) + assert ((result[[0, 1, 2, 3]] == data[indexer]).all()) + assert (result.dtype == dtype) + + def test_2d_with_out(self, dtype_can_hold_na, writeable): + dtype, can_hold_na = dtype_can_hold_na + + data = np.random.randint(0, 2, (5, 3)).astype(dtype) + data.flags.writeable = writeable + + indexer = [2, 1, 0, 1] + out0 = np.empty((4, 3), dtype=dtype) + out1 = np.empty((5, 4), dtype=dtype) + algos.take_nd(data, indexer, out=out0, axis=0) + algos.take_nd(data, indexer, out=out1, axis=1) + + expected0 = data.take(indexer, axis=0) + expected1 = data.take(indexer, axis=1) + tm.assert_almost_equal(out0, expected0) + tm.assert_almost_equal(out1, expected1) + + indexer = [2, 1, 0, -1] + out0 = np.empty((4, 3), dtype=dtype) + out1 = np.empty((5, 4), dtype=dtype) + + if can_hold_na: algos.take_nd(data, indexer, out=out0, axis=0) algos.take_nd(data, indexer, out=out1, axis=1) + expected0 = data.take(indexer, axis=0) expected1 = data.take(indexer, axis=1) + expected0[3, :] = np.nan + expected1[:, 3] = np.nan + tm.assert_almost_equal(out0, expected0) tm.assert_almost_equal(out1, expected1) - - indexer = [2, 1, 0, -1] - out0 = np.empty((4, 3), dtype=dtype) - out1 = np.empty((5, 4), dtype=dtype) - if can_hold_na: - algos.take_nd(data, indexer, out=out0, axis=0) - algos.take_nd(data, indexer, out=out1, axis=1) - expected0 = data.take(indexer, axis=0) - expected1 = data.take(indexer, axis=1) - expected0[3, :] = np.nan - expected1[:, 3] = np.nan - tm.assert_almost_equal(out0, expected0) - tm.assert_almost_equal(out1, expected1) - else: - for i, out in enumerate([out0, out1]): - with tm.assert_raises_regex(TypeError, - self.fill_error): - algos.take_nd(data, indexer, out=out, axis=i) - # no exception o/w - data.take(indexer, out=out, axis=i) - - for writeable in [True, False]: - # Check that take_nd works both with writeable arrays (in which - # case fast typed memoryviews implementation) and read-only - # arrays alike. - _test_dtype(np.float64, True, writeable=writeable) - _test_dtype(np.float32, True, writeable=writeable) - _test_dtype(np.uint64, False, writeable=writeable) - _test_dtype(np.uint32, False, writeable=writeable) - _test_dtype(np.uint16, False, writeable=writeable) - _test_dtype(np.uint8, False, writeable=writeable) - _test_dtype(np.int64, False, writeable=writeable) - _test_dtype(np.int32, False, writeable=writeable) - _test_dtype(np.int16, False, writeable=writeable) - _test_dtype(np.int8, False, writeable=writeable) - _test_dtype(np.object_, True, writeable=writeable) - _test_dtype(np.bool, False, writeable=writeable) - - def test_2d_fill_nonna(self): - def _test_dtype(dtype, fill_value, out_dtype): - data = np.random.randint(0, 2, (5, 3)).astype(dtype) - - indexer = [2, 1, 0, -1] - - result = algos.take_nd(data, indexer, axis=0, - fill_value=fill_value) - assert ((result[[0, 1, 2], :] == data[[2, 1, 0], :]).all()) - assert ((result[3, :] == fill_value).all()) - assert (result.dtype == out_dtype) - - result = algos.take_nd(data, indexer, axis=1, - fill_value=fill_value) - assert ((result[:, [0, 1, 2]] == data[:, [2, 1, 0]]).all()) - assert ((result[:, 3] == fill_value).all()) - assert (result.dtype == out_dtype) - - indexer = [2, 1, 0, 1] - - result = algos.take_nd(data, indexer, axis=0, - fill_value=fill_value) - assert ((result[[0, 1, 2, 3], :] == data[indexer, :]).all()) - assert (result.dtype == dtype) - - result = algos.take_nd(data, indexer, axis=1, - fill_value=fill_value) - assert ((result[:, [0, 1, 2, 3]] == data[:, indexer]).all()) - assert (result.dtype == dtype) - - _test_dtype(np.int8, np.int16(127), np.int8) - _test_dtype(np.int8, np.int16(128), np.int16) - _test_dtype(np.int32, 1, np.int32) - _test_dtype(np.int32, 2.0, np.float64) - _test_dtype(np.int32, 3.0 + 4.0j, np.complex128) - _test_dtype(np.int32, True, np.object_) - _test_dtype(np.int32, '', np.object_) - _test_dtype(np.float64, 1, np.float64) - _test_dtype(np.float64, 2.0, np.float64) - _test_dtype(np.float64, 3.0 + 4.0j, np.complex128) - _test_dtype(np.float64, True, np.object_) - _test_dtype(np.float64, '', np.object_) - _test_dtype(np.complex128, 1, np.complex128) - _test_dtype(np.complex128, 2.0, np.complex128) - _test_dtype(np.complex128, 3.0 + 4.0j, np.complex128) - _test_dtype(np.complex128, True, np.object_) - _test_dtype(np.complex128, '', np.object_) - _test_dtype(np.bool_, 1, np.object_) - _test_dtype(np.bool_, 2.0, np.object_) - _test_dtype(np.bool_, 3.0 + 4.0j, np.object_) - _test_dtype(np.bool_, True, np.bool_) - _test_dtype(np.bool_, '', np.object_) - - def test_3d_with_out(self): - def _test_dtype(dtype, can_hold_na): - data = np.random.randint(0, 2, (5, 4, 3)).astype(dtype) - - indexer = [2, 1, 0, 1] - out0 = np.empty((4, 4, 3), dtype=dtype) - out1 = np.empty((5, 4, 3), dtype=dtype) - out2 = np.empty((5, 4, 4), dtype=dtype) + else: + for i, out in enumerate([out0, out1]): + with tm.assert_raises_regex(TypeError, + self.fill_error): + algos.take_nd(data, indexer, out=out, axis=i) + + # No Exception otherwise. + data.take(indexer, out=out, axis=i) + + def test_2d_fill_nonna(self, dtype_fill_out_dtype): + dtype, fill_value, out_dtype = dtype_fill_out_dtype + data = np.random.randint(0, 2, (5, 3)).astype(dtype) + indexer = [2, 1, 0, -1] + + result = algos.take_nd(data, indexer, axis=0, + fill_value=fill_value) + assert ((result[[0, 1, 2], :] == data[[2, 1, 0], :]).all()) + assert ((result[3, :] == fill_value).all()) + assert (result.dtype == out_dtype) + + result = algos.take_nd(data, indexer, axis=1, + fill_value=fill_value) + assert ((result[:, [0, 1, 2]] == data[:, [2, 1, 0]]).all()) + assert ((result[:, 3] == fill_value).all()) + assert (result.dtype == out_dtype) + + indexer = [2, 1, 0, 1] + result = algos.take_nd(data, indexer, axis=0, + fill_value=fill_value) + assert ((result[[0, 1, 2, 3], :] == data[indexer, :]).all()) + assert (result.dtype == dtype) + + result = algos.take_nd(data, indexer, axis=1, + fill_value=fill_value) + assert ((result[:, [0, 1, 2, 3]] == data[:, indexer]).all()) + assert (result.dtype == dtype) + + def test_3d_with_out(self, dtype_can_hold_na): + dtype, can_hold_na = dtype_can_hold_na + + data = np.random.randint(0, 2, (5, 4, 3)).astype(dtype) + indexer = [2, 1, 0, 1] + + out0 = np.empty((4, 4, 3), dtype=dtype) + out1 = np.empty((5, 4, 3), dtype=dtype) + out2 = np.empty((5, 4, 4), dtype=dtype) + + algos.take_nd(data, indexer, out=out0, axis=0) + algos.take_nd(data, indexer, out=out1, axis=1) + algos.take_nd(data, indexer, out=out2, axis=2) + + expected0 = data.take(indexer, axis=0) + expected1 = data.take(indexer, axis=1) + expected2 = data.take(indexer, axis=2) + + tm.assert_almost_equal(out0, expected0) + tm.assert_almost_equal(out1, expected1) + tm.assert_almost_equal(out2, expected2) + + indexer = [2, 1, 0, -1] + out0 = np.empty((4, 4, 3), dtype=dtype) + out1 = np.empty((5, 4, 3), dtype=dtype) + out2 = np.empty((5, 4, 4), dtype=dtype) + + if can_hold_na: algos.take_nd(data, indexer, out=out0, axis=0) algos.take_nd(data, indexer, out=out1, axis=1) algos.take_nd(data, indexer, out=out2, axis=2) + expected0 = data.take(indexer, axis=0) expected1 = data.take(indexer, axis=1) expected2 = data.take(indexer, axis=2) + + expected0[3, :, :] = np.nan + expected1[:, 3, :] = np.nan + expected2[:, :, 3] = np.nan + tm.assert_almost_equal(out0, expected0) tm.assert_almost_equal(out1, expected1) tm.assert_almost_equal(out2, expected2) - - indexer = [2, 1, 0, -1] - out0 = np.empty((4, 4, 3), dtype=dtype) - out1 = np.empty((5, 4, 3), dtype=dtype) - out2 = np.empty((5, 4, 4), dtype=dtype) - if can_hold_na: - algos.take_nd(data, indexer, out=out0, axis=0) - algos.take_nd(data, indexer, out=out1, axis=1) - algos.take_nd(data, indexer, out=out2, axis=2) - expected0 = data.take(indexer, axis=0) - expected1 = data.take(indexer, axis=1) - expected2 = data.take(indexer, axis=2) - expected0[3, :, :] = np.nan - expected1[:, 3, :] = np.nan - expected2[:, :, 3] = np.nan - tm.assert_almost_equal(out0, expected0) - tm.assert_almost_equal(out1, expected1) - tm.assert_almost_equal(out2, expected2) - else: - for i, out in enumerate([out0, out1, out2]): - with tm.assert_raises_regex(TypeError, - self.fill_error): - algos.take_nd(data, indexer, out=out, axis=i) - # no exception o/w - data.take(indexer, out=out, axis=i) - - _test_dtype(np.float64, True) - _test_dtype(np.float32, True) - _test_dtype(np.uint64, False) - _test_dtype(np.uint32, False) - _test_dtype(np.uint16, False) - _test_dtype(np.uint8, False) - _test_dtype(np.int64, False) - _test_dtype(np.int32, False) - _test_dtype(np.int16, False) - _test_dtype(np.int8, False) - _test_dtype(np.object_, True) - _test_dtype(np.bool, False) - - def test_3d_fill_nonna(self): - def _test_dtype(dtype, fill_value, out_dtype): - data = np.random.randint(0, 2, (5, 4, 3)).astype(dtype) - - indexer = [2, 1, 0, -1] - - result = algos.take_nd(data, indexer, axis=0, - fill_value=fill_value) - assert ((result[[0, 1, 2], :, :] == data[[2, 1, 0], :, :]).all()) - assert ((result[3, :, :] == fill_value).all()) - assert (result.dtype == out_dtype) - - result = algos.take_nd(data, indexer, axis=1, - fill_value=fill_value) - assert ((result[:, [0, 1, 2], :] == data[:, [2, 1, 0], :]).all()) - assert ((result[:, 3, :] == fill_value).all()) - assert (result.dtype == out_dtype) - - result = algos.take_nd(data, indexer, axis=2, - fill_value=fill_value) - assert ((result[:, :, [0, 1, 2]] == data[:, :, [2, 1, 0]]).all()) - assert ((result[:, :, 3] == fill_value).all()) - assert (result.dtype == out_dtype) - - indexer = [2, 1, 0, 1] - - result = algos.take_nd(data, indexer, axis=0, - fill_value=fill_value) - assert ((result[[0, 1, 2, 3], :, :] == data[indexer, :, :]).all()) - assert (result.dtype == dtype) - - result = algos.take_nd(data, indexer, axis=1, - fill_value=fill_value) - assert ((result[:, [0, 1, 2, 3], :] == data[:, indexer, :]).all()) - assert (result.dtype == dtype) - - result = algos.take_nd(data, indexer, axis=2, - fill_value=fill_value) - assert ((result[:, :, [0, 1, 2, 3]] == data[:, :, indexer]).all()) - assert (result.dtype == dtype) - - _test_dtype(np.int8, np.int16(127), np.int8) - _test_dtype(np.int8, np.int16(128), np.int16) - _test_dtype(np.int32, 1, np.int32) - _test_dtype(np.int32, 2.0, np.float64) - _test_dtype(np.int32, 3.0 + 4.0j, np.complex128) - _test_dtype(np.int32, True, np.object_) - _test_dtype(np.int32, '', np.object_) - _test_dtype(np.float64, 1, np.float64) - _test_dtype(np.float64, 2.0, np.float64) - _test_dtype(np.float64, 3.0 + 4.0j, np.complex128) - _test_dtype(np.float64, True, np.object_) - _test_dtype(np.float64, '', np.object_) - _test_dtype(np.complex128, 1, np.complex128) - _test_dtype(np.complex128, 2.0, np.complex128) - _test_dtype(np.complex128, 3.0 + 4.0j, np.complex128) - _test_dtype(np.complex128, True, np.object_) - _test_dtype(np.complex128, '', np.object_) - _test_dtype(np.bool_, 1, np.object_) - _test_dtype(np.bool_, 2.0, np.object_) - _test_dtype(np.bool_, 3.0 + 4.0j, np.object_) - _test_dtype(np.bool_, True, np.bool_) - _test_dtype(np.bool_, '', np.object_) + else: + for i, out in enumerate([out0, out1, out2]): + with tm.assert_raises_regex(TypeError, + self.fill_error): + algos.take_nd(data, indexer, out=out, axis=i) + + # No Exception otherwise. + data.take(indexer, out=out, axis=i) + + def test_3d_fill_nonna(self, dtype_fill_out_dtype): + dtype, fill_value, out_dtype = dtype_fill_out_dtype + + data = np.random.randint(0, 2, (5, 4, 3)).astype(dtype) + indexer = [2, 1, 0, -1] + + result = algos.take_nd(data, indexer, axis=0, + fill_value=fill_value) + assert ((result[[0, 1, 2], :, :] == data[[2, 1, 0], :, :]).all()) + assert ((result[3, :, :] == fill_value).all()) + assert (result.dtype == out_dtype) + + result = algos.take_nd(data, indexer, axis=1, + fill_value=fill_value) + assert ((result[:, [0, 1, 2], :] == data[:, [2, 1, 0], :]).all()) + assert ((result[:, 3, :] == fill_value).all()) + assert (result.dtype == out_dtype) + + result = algos.take_nd(data, indexer, axis=2, + fill_value=fill_value) + assert ((result[:, :, [0, 1, 2]] == data[:, :, [2, 1, 0]]).all()) + assert ((result[:, :, 3] == fill_value).all()) + assert (result.dtype == out_dtype) + + indexer = [2, 1, 0, 1] + result = algos.take_nd(data, indexer, axis=0, + fill_value=fill_value) + assert ((result[[0, 1, 2, 3], :, :] == data[indexer, :, :]).all()) + assert (result.dtype == dtype) + + result = algos.take_nd(data, indexer, axis=1, + fill_value=fill_value) + assert ((result[:, [0, 1, 2, 3], :] == data[:, indexer, :]).all()) + assert (result.dtype == dtype) + + result = algos.take_nd(data, indexer, axis=2, + fill_value=fill_value) + assert ((result[:, :, [0, 1, 2, 3]] == data[:, :, indexer]).all()) + assert (result.dtype == dtype) def test_1d_other_dtypes(self): arr = np.random.randn(10).astype(np.float32)