diff --git a/pandas/tests/indexing/common.py b/pandas/tests/indexing/common.py index ea9f2584196d3..2af76f69a4300 100644 --- a/pandas/tests/indexing/common.py +++ b/pandas/tests/indexing/common.py @@ -1,190 +1,40 @@ """ common utilities """ -import itertools +from __future__ import annotations -import numpy as np - -from pandas import ( - DataFrame, - MultiIndex, - Series, - date_range, -) -import pandas._testing as tm -from pandas.core.api import ( - Float64Index, - UInt64Index, +from typing import ( + Any, + Literal, ) -def _mklbl(prefix, n): +def _mklbl(prefix: str, n: int): return [f"{prefix}{i}" for i in range(n)] -def _axify(obj, key, axis): - # create a tuple accessor - axes = [slice(None)] * obj.ndim - axes[axis] = key - return tuple(axes) - - -class Base: - """indexing comprehensive base class""" - - _kinds = {"series", "frame"} - _typs = { - "ints", - "uints", - "labels", - "mixed", - "ts", - "floats", - "empty", - "ts_rev", - "multi", - } - - def setup_method(self): - - self.series_ints = Series(np.random.rand(4), index=np.arange(0, 8, 2)) - self.frame_ints = DataFrame( - np.random.randn(4, 4), index=np.arange(0, 8, 2), columns=np.arange(0, 12, 3) - ) - - self.series_uints = Series( - np.random.rand(4), index=UInt64Index(np.arange(0, 8, 2)) - ) - self.frame_uints = DataFrame( - np.random.randn(4, 4), - index=UInt64Index(range(0, 8, 2)), - columns=UInt64Index(range(0, 12, 3)), - ) - - self.series_floats = Series( - np.random.rand(4), index=Float64Index(range(0, 8, 2)) - ) - self.frame_floats = DataFrame( - np.random.randn(4, 4), - index=Float64Index(range(0, 8, 2)), - columns=Float64Index(range(0, 12, 3)), - ) - - m_idces = [ - MultiIndex.from_product([[1, 2], [3, 4]]), - MultiIndex.from_product([[5, 6], [7, 8]]), - MultiIndex.from_product([[9, 10], [11, 12]]), - ] - - self.series_multi = Series(np.random.rand(4), index=m_idces[0]) - self.frame_multi = DataFrame( - np.random.randn(4, 4), index=m_idces[0], columns=m_idces[1] - ) - - self.series_labels = Series(np.random.randn(4), index=list("abcd")) - self.frame_labels = DataFrame( - np.random.randn(4, 4), index=list("abcd"), columns=list("ABCD") - ) - - self.series_mixed = Series(np.random.randn(4), index=[2, 4, "null", 8]) - self.frame_mixed = DataFrame(np.random.randn(4, 4), index=[2, 4, "null", 8]) - - self.series_ts = Series( - np.random.randn(4), index=date_range("20130101", periods=4) - ) - self.frame_ts = DataFrame( - np.random.randn(4, 4), index=date_range("20130101", periods=4) - ) - - dates_rev = date_range("20130101", periods=4).sort_values(ascending=False) - self.series_ts_rev = Series(np.random.randn(4), index=dates_rev) - self.frame_ts_rev = DataFrame(np.random.randn(4, 4), index=dates_rev) - - self.frame_empty = DataFrame() - self.series_empty = Series(dtype=object) - - # form agglomerates - for kind in self._kinds: - d = {} - for typ in self._typs: - d[typ] = getattr(self, f"{kind}_{typ}") - - setattr(self, kind, d) - - def generate_indices(self, f, values=False): - """ - generate the indices - if values is True , use the axis values - is False, use the range - """ - axes = f.axes - if values: - axes = (list(range(len(ax))) for ax in axes) - - return itertools.product(*axes) - - def get_value(self, name, f, i, values=False): - """return the value for the location i""" - # check against values - if values: - return f.values[i] - - elif name == "iat": - return f.iloc[i] - else: - assert name == "at" - return f.loc[i] - - def check_values(self, f, func, values=False): - - if f is None: - return - axes = f.axes - indices = itertools.product(*axes) - - for i in indices: - result = getattr(f, func)[i] - - # check against values - if values: - expected = f.values[i] - else: - expected = f - for a in reversed(i): - expected = expected.__getitem__(a) - - tm.assert_almost_equal(result, expected) - - def check_result(self, method, key, typs=None, axes=None, fails=None): - def _eq(axis, obj, key): - """compare equal for these 2 keys""" - axified = _axify(obj, key, axis) +def check_indexing_smoketest_or_raises( + obj, + method: Literal["iloc", "loc"], + key: Any, + axes: Literal[0, 1] | None = None, + fails=None, +) -> None: + if axes is None: + axes_list = [0, 1] + else: + assert axes in [0, 1] + axes_list = [axes] + + for ax in axes_list: + if ax < obj.ndim: + # create a tuple accessor + new_axes = [slice(None)] * obj.ndim + new_axes[ax] = key + axified = tuple(new_axes) try: getattr(obj, method).__getitem__(axified) - except (IndexError, TypeError, KeyError) as detail: - # if we are in fails, the ok, otherwise raise it if fails is not None: if isinstance(detail, fails): return raise - - if typs is None: - typs = self._typs - - if axes is None: - axes = [0, 1] - else: - assert axes in [0, 1] - axes = [axes] - - # check - for kind in self._kinds: - - d = getattr(self, kind) - for ax in axes: - for typ in typs: - assert typ in self._typs - - obj = d[typ] - if ax < obj.ndim: - _eq(axis=ax, obj=obj, key=key) diff --git a/pandas/tests/indexing/conftest.py b/pandas/tests/indexing/conftest.py new file mode 100644 index 0000000000000..ac3db524170fb --- /dev/null +++ b/pandas/tests/indexing/conftest.py @@ -0,0 +1,107 @@ +import numpy as np +import pytest + +from pandas import ( + DataFrame, + MultiIndex, + Series, + date_range, +) +from pandas.core.api import ( + Float64Index, + UInt64Index, +) + + +@pytest.fixture +def series_ints(): + return Series(np.random.rand(4), index=np.arange(0, 8, 2)) + + +@pytest.fixture +def frame_ints(): + return DataFrame( + np.random.randn(4, 4), index=np.arange(0, 8, 2), columns=np.arange(0, 12, 3) + ) + + +@pytest.fixture +def series_uints(): + return Series(np.random.rand(4), index=UInt64Index(np.arange(0, 8, 2))) + + +@pytest.fixture +def frame_uints(): + return DataFrame( + np.random.randn(4, 4), + index=UInt64Index(range(0, 8, 2)), + columns=UInt64Index(range(0, 12, 3)), + ) + + +@pytest.fixture +def series_labels(): + return Series(np.random.randn(4), index=list("abcd")) + + +@pytest.fixture +def frame_labels(): + return DataFrame(np.random.randn(4, 4), index=list("abcd"), columns=list("ABCD")) + + +@pytest.fixture +def series_ts(): + return Series(np.random.randn(4), index=date_range("20130101", periods=4)) + + +@pytest.fixture +def frame_ts(): + return DataFrame(np.random.randn(4, 4), index=date_range("20130101", periods=4)) + + +@pytest.fixture +def series_floats(): + return Series(np.random.rand(4), index=Float64Index(range(0, 8, 2))) + + +@pytest.fixture +def frame_floats(): + return DataFrame( + np.random.randn(4, 4), + index=Float64Index(range(0, 8, 2)), + columns=Float64Index(range(0, 12, 3)), + ) + + +@pytest.fixture +def series_mixed(): + return Series(np.random.randn(4), index=[2, 4, "null", 8]) + + +@pytest.fixture +def frame_mixed(): + return DataFrame(np.random.randn(4, 4), index=[2, 4, "null", 8]) + + +@pytest.fixture +def frame_empty(): + return DataFrame() + + +@pytest.fixture +def series_empty(): + return Series(dtype=object) + + +@pytest.fixture +def frame_multi(): + return DataFrame( + np.random.randn(4, 4), + index=MultiIndex.from_product([[1, 2], [3, 4]]), + columns=MultiIndex.from_product([[5, 6], [7, 8]]), + ) + + +@pytest.fixture +def series_multi(): + return Series(np.random.rand(4), index=MultiIndex.from_product([[1, 2], [3, 4]])) diff --git a/pandas/tests/indexing/test_iloc.py b/pandas/tests/indexing/test_iloc.py index 8cc6b6e73aaea..db2fe45faf6de 100644 --- a/pandas/tests/indexing/test_iloc.py +++ b/pandas/tests/indexing/test_iloc.py @@ -32,7 +32,7 @@ ) import pandas._testing as tm from pandas.api.types import is_scalar -from pandas.tests.indexing.common import Base +from pandas.tests.indexing.common import check_indexing_smoketest_or_raises # We pass through the error message from numpy _slice_iloc_msg = re.escape( @@ -41,13 +41,19 @@ ) -class TestiLoc(Base): +class TestiLoc: @pytest.mark.parametrize("key", [2, -1, [0, 1, 2]]) - def test_iloc_getitem_int_and_list_int(self, key): - self.check_result( + @pytest.mark.parametrize("kind", ["series", "frame"]) + @pytest.mark.parametrize( + "col", + ["labels", "mixed", "ts", "floats", "empty"], + ) + def test_iloc_getitem_int_and_list_int(self, key, kind, col, request): + obj = request.getfixturevalue(f"{kind}_{col}") + check_indexing_smoketest_or_raises( + obj, "iloc", key, - typs=["labels", "mixed", "ts", "floats", "empty"], fails=IndexError, ) diff --git a/pandas/tests/indexing/test_loc.py b/pandas/tests/indexing/test_loc.py index 4d451c634fb09..3490d05f13e9d 100644 --- a/pandas/tests/indexing/test_loc.py +++ b/pandas/tests/indexing/test_loc.py @@ -40,7 +40,7 @@ from pandas.api.types import is_scalar from pandas.core.api import Float64Index from pandas.core.indexing import _one_ellipsis_message -from pandas.tests.indexing.common import Base +from pandas.tests.indexing.common import check_indexing_smoketest_or_raises @pytest.mark.parametrize( @@ -59,16 +59,20 @@ def test_not_change_nan_loc(series, new_series, expected_ser): tm.assert_frame_equal(df.notna(), ~expected) -class TestLoc(Base): - def test_loc_getitem_int(self): +class TestLoc: + @pytest.mark.parametrize("kind", ["series", "frame"]) + def test_loc_getitem_int(self, kind, request): # int label - self.check_result("loc", 2, typs=["labels"], fails=KeyError) + obj = request.getfixturevalue(f"{kind}_labels") + check_indexing_smoketest_or_raises(obj, "loc", 2, fails=KeyError) - def test_loc_getitem_label(self): + @pytest.mark.parametrize("kind", ["series", "frame"]) + def test_loc_getitem_label(self, kind, request): # label - self.check_result("loc", "c", typs=["empty"], fails=KeyError) + obj = request.getfixturevalue(f"{kind}_empty") + check_indexing_smoketest_or_raises(obj, "loc", "c", fails=KeyError) @pytest.mark.parametrize( "key, typs, axes", @@ -81,10 +85,14 @@ def test_loc_getitem_label(self): [20, ["floats"], 0], ], ) - def test_loc_getitem_label_out_of_range(self, key, typs, axes): - - # out of range label - self.check_result("loc", key, typs=typs, axes=axes, fails=KeyError) + @pytest.mark.parametrize("kind", ["series", "frame"]) + def test_loc_getitem_label_out_of_range(self, key, typs, axes, kind, request): + for typ in typs: + obj = request.getfixturevalue(f"{kind}_{typ}") + # out of range label + check_indexing_smoketest_or_raises( + obj, "loc", key, axes=axes, fails=KeyError + ) @pytest.mark.parametrize( "key, typs", @@ -93,9 +101,12 @@ def test_loc_getitem_label_out_of_range(self, key, typs, axes): [[1, 3.0, "A"], ["ints", "uints", "floats"]], ], ) - def test_loc_getitem_label_list(self, key, typs): - # list of labels - self.check_result("loc", key, typs=typs, fails=KeyError) + @pytest.mark.parametrize("kind", ["series", "frame"]) + def test_loc_getitem_label_list(self, key, typs, kind, request): + for typ in typs: + obj = request.getfixturevalue(f"{kind}_{typ}") + # list of labels + check_indexing_smoketest_or_raises(obj, "loc", key, fails=KeyError) @pytest.mark.parametrize( "key, typs, axes", @@ -107,13 +118,21 @@ def test_loc_getitem_label_list(self, key, typs): [[(1, 3), (1, 4), (2, 5)], ["multi"], 0], ], ) - def test_loc_getitem_label_list_with_missing(self, key, typs, axes): - self.check_result("loc", key, typs=typs, axes=axes, fails=KeyError) + @pytest.mark.parametrize("kind", ["series", "frame"]) + def test_loc_getitem_label_list_with_missing(self, key, typs, axes, kind, request): + for typ in typs: + obj = request.getfixturevalue(f"{kind}_{typ}") + check_indexing_smoketest_or_raises( + obj, "loc", key, axes=axes, fails=KeyError + ) - def test_loc_getitem_label_list_fails(self): + @pytest.mark.parametrize("typs", ["ints", "uints"]) + @pytest.mark.parametrize("kind", ["series", "frame"]) + def test_loc_getitem_label_list_fails(self, typs, kind, request): # fails - self.check_result( - "loc", [20, 30, 40], typs=["ints", "uints"], axes=1, fails=KeyError + obj = request.getfixturevalue(f"{kind}_{typs}") + check_indexing_smoketest_or_raises( + obj, "loc", [20, 30, 40], axes=1, fails=KeyError ) def test_loc_getitem_label_array_like(self): @@ -121,11 +140,13 @@ def test_loc_getitem_label_array_like(self): # array like pass - def test_loc_getitem_bool(self): + @pytest.mark.parametrize("kind", ["series", "frame"]) + def test_loc_getitem_bool(self, kind, request): + obj = request.getfixturevalue(f"{kind}_empty") # boolean indexers b = [True, False, True, False] - self.check_result("loc", b, typs=["empty"], fails=IndexError) + check_indexing_smoketest_or_raises(obj, "loc", b, fails=IndexError) @pytest.mark.parametrize( "slc, typs, axes, fails", @@ -142,21 +163,23 @@ def test_loc_getitem_bool(self): [slice(2, 4, 2), ["mixed"], 0, TypeError], ], ) - def test_loc_getitem_label_slice(self, slc, typs, axes, fails): + @pytest.mark.parametrize("kind", ["series", "frame"]) + def test_loc_getitem_label_slice(self, slc, typs, axes, fails, kind, request): # label slices (with ints) # real label slices # GH 14316 - - self.check_result( - "loc", - slc, - typs=typs, - axes=axes, - fails=fails, - ) + for typ in typs: + obj = request.getfixturevalue(f"{kind}_{typ}") + check_indexing_smoketest_or_raises( + obj, + "loc", + slc, + axes=axes, + fails=fails, + ) def test_setitem_from_duplicate_axis(self): # GH#34034 @@ -2465,7 +2488,7 @@ def test_loc_getitem_slice_columns_mixed_dtype(self): tm.assert_frame_equal(df.loc[:, 1:], expected) -class TestLocBooleanLabelsAndSlices(Base): +class TestLocBooleanLabelsAndSlices: @pytest.mark.parametrize("bool_value", [True, False]) def test_loc_bool_incompatible_index_raises( self, index, frame_or_series, bool_value diff --git a/pandas/tests/indexing/test_scalar.py b/pandas/tests/indexing/test_scalar.py index 552db9c5e164d..b1c3008b04d3e 100644 --- a/pandas/tests/indexing/test_scalar.py +++ b/pandas/tests/indexing/test_scalar.py @@ -3,6 +3,7 @@ datetime, timedelta, ) +import itertools import numpy as np import pytest @@ -15,44 +16,50 @@ date_range, ) import pandas._testing as tm -from pandas.tests.indexing.common import Base -class TestScalar(Base): +def generate_indices(f, values=False): + """ + generate the indices + if values is True , use the axis values + is False, use the range + """ + axes = f.axes + if values: + axes = (list(range(len(ax))) for ax in axes) + + return itertools.product(*axes) + + +class TestScalar: @pytest.mark.parametrize("kind", ["series", "frame"]) @pytest.mark.parametrize("col", ["ints", "uints"]) - def test_iat_set_ints(self, kind, col): - f = getattr(self, kind)[col] - if f is not None: - indices = self.generate_indices(f, True) - for i in indices: - f.iat[i] = 1 - expected = self.get_value("iat", f, i, True) - tm.assert_almost_equal(expected, 1) + def test_iat_set_ints(self, kind, col, request): + f = request.getfixturevalue(f"{kind}_{col}") + indices = generate_indices(f, True) + for i in indices: + f.iat[i] = 1 + expected = f.values[i] + tm.assert_almost_equal(expected, 1) @pytest.mark.parametrize("kind", ["series", "frame"]) @pytest.mark.parametrize("col", ["labels", "ts", "floats"]) - def test_iat_set_other(self, kind, col): - f = getattr(self, kind)[col] - if f is not None: - msg = "iAt based indexing can only have integer indexers" - with pytest.raises(ValueError, match=msg): - indices = self.generate_indices(f, False) - for i in indices: - f.iat[i] = 1 - expected = self.get_value("iat", f, i, False) - tm.assert_almost_equal(expected, 1) + def test_iat_set_other(self, kind, col, request): + f = request.getfixturevalue(f"{kind}_{col}") + msg = "iAt based indexing can only have integer indexers" + with pytest.raises(ValueError, match=msg): + idx = next(generate_indices(f, False)) + f.iat[idx] = 1 @pytest.mark.parametrize("kind", ["series", "frame"]) @pytest.mark.parametrize("col", ["ints", "uints", "labels", "ts", "floats"]) - def test_at_set_ints_other(self, kind, col): - f = getattr(self, kind)[col] - if f is not None: - indices = self.generate_indices(f, False) - for i in indices: - f.at[i] = 1 - expected = self.get_value("at", f, i, False) - tm.assert_almost_equal(expected, 1) + def test_at_set_ints_other(self, kind, col, request): + f = request.getfixturevalue(f"{kind}_{col}") + indices = generate_indices(f, False) + for i in indices: + f.at[i] = 1 + expected = f.loc[i] + tm.assert_almost_equal(expected, 1) class TestAtAndiAT: