| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199 |
- import numpy as np
- import pytest
- import pandas as pd
- from pandas import Index
- import pandas._testing as tm
- def _isnan(val):
- try:
- return val is not pd.NA and np.isnan(val)
- except TypeError:
- return False
- def _equivalent_na(dtype, null):
- if dtype.na_value is pd.NA and null is pd.NA:
- return True
- elif _isnan(dtype.na_value) and _isnan(null):
- return True
- else:
- return False
- class TestGetLoc:
- def test_get_loc(self, any_string_dtype):
- index = Index(["a", "b", "c"], dtype=any_string_dtype)
- assert index.get_loc("b") == 1
- def test_get_loc_raises(self, any_string_dtype):
- index = Index(["a", "b", "c"], dtype=any_string_dtype)
- with pytest.raises(KeyError, match="d"):
- index.get_loc("d")
- def test_get_loc_invalid_value(self, any_string_dtype):
- index = Index(["a", "b", "c"], dtype=any_string_dtype)
- with pytest.raises(KeyError, match="1"):
- index.get_loc(1)
- def test_get_loc_non_unique(self, any_string_dtype):
- index = Index(["a", "b", "a"], dtype=any_string_dtype)
- result = index.get_loc("a")
- expected = np.array([True, False, True])
- tm.assert_numpy_array_equal(result, expected)
- def test_get_loc_non_missing(self, any_string_dtype, nulls_fixture):
- index = Index(["a", "b", "c"], dtype=any_string_dtype)
- with pytest.raises(KeyError):
- index.get_loc(nulls_fixture)
- def test_get_loc_missing(self, any_string_dtype, nulls_fixture):
- index = Index(["a", "b", nulls_fixture], dtype=any_string_dtype)
- assert index.get_loc(nulls_fixture) == 2
- class TestGetIndexer:
- @pytest.mark.parametrize(
- "method,expected",
- [
- ("pad", [-1, 0, 1, 1]),
- ("backfill", [0, 0, 1, -1]),
- ],
- )
- def test_get_indexer_strings(self, any_string_dtype, method, expected):
- expected = np.array(expected, dtype=np.intp)
- index = Index(["b", "c"], dtype=any_string_dtype)
- actual = index.get_indexer(["a", "b", "c", "d"], method=method)
- tm.assert_numpy_array_equal(actual, expected)
- def test_get_indexer_strings_raises(self, any_string_dtype):
- index = Index(["b", "c"], dtype=any_string_dtype)
- msg = "|".join(
- [
- "operation 'sub' not supported for dtype 'str",
- r"unsupported operand type\(s\) for -: 'str' and 'str'",
- ]
- )
- with pytest.raises(TypeError, match=msg):
- index.get_indexer(["a", "b", "c", "d"], method="nearest")
- with pytest.raises(TypeError, match=msg):
- index.get_indexer(["a", "b", "c", "d"], method="pad", tolerance=2)
- with pytest.raises(TypeError, match=msg):
- index.get_indexer(
- ["a", "b", "c", "d"], method="pad", tolerance=[2, 2, 2, 2]
- )
- @pytest.mark.parametrize("null", [None, np.nan, float("nan"), pd.NA])
- def test_get_indexer_missing(self, any_string_dtype, null, using_infer_string):
- # NaT and Decimal("NaN") from null_fixture are not supported for string dtype
- index = Index(["a", "b", null], dtype=any_string_dtype)
- result = index.get_indexer(["a", null, "c"])
- if using_infer_string:
- expected = np.array([0, 2, -1], dtype=np.intp)
- elif any_string_dtype == "string" and not _equivalent_na(
- any_string_dtype, null
- ):
- expected = np.array([0, -1, -1], dtype=np.intp)
- else:
- expected = np.array([0, 2, -1], dtype=np.intp)
- tm.assert_numpy_array_equal(result, expected)
- class TestGetIndexerNonUnique:
- @pytest.mark.parametrize("null", [None, np.nan, float("nan"), pd.NA])
- def test_get_indexer_non_unique_nas(
- self, any_string_dtype, null, using_infer_string
- ):
- index = Index(["a", "b", null], dtype=any_string_dtype)
- indexer, missing = index.get_indexer_non_unique(["a", null])
- if using_infer_string:
- expected_indexer = np.array([0, 2], dtype=np.intp)
- expected_missing = np.array([], dtype=np.intp)
- elif any_string_dtype == "string" and not _equivalent_na(
- any_string_dtype, null
- ):
- expected_indexer = np.array([0, -1], dtype=np.intp)
- expected_missing = np.array([1], dtype=np.intp)
- else:
- expected_indexer = np.array([0, 2], dtype=np.intp)
- expected_missing = np.array([], dtype=np.intp)
- tm.assert_numpy_array_equal(indexer, expected_indexer)
- tm.assert_numpy_array_equal(missing, expected_missing)
- # actually non-unique
- index = Index(["a", null, "b", null], dtype=any_string_dtype)
- indexer, missing = index.get_indexer_non_unique(["a", null])
- if using_infer_string:
- expected_indexer = np.array([0, 1, 3], dtype=np.intp)
- elif any_string_dtype == "string" and not _equivalent_na(
- any_string_dtype, null
- ):
- pass
- else:
- expected_indexer = np.array([0, 1, 3], dtype=np.intp)
- tm.assert_numpy_array_equal(indexer, expected_indexer)
- tm.assert_numpy_array_equal(missing, expected_missing)
- class TestSliceLocs:
- @pytest.mark.parametrize(
- "in_slice,expected",
- [
- # error: Slice index must be an integer or None
- (pd.IndexSlice[::-1], "yxdcb"),
- (pd.IndexSlice["b":"y":-1], ""), # type: ignore[misc]
- (pd.IndexSlice["b"::-1], "b"), # type: ignore[misc]
- (pd.IndexSlice[:"b":-1], "yxdcb"), # type: ignore[misc]
- (pd.IndexSlice[:"y":-1], "y"), # type: ignore[misc]
- (pd.IndexSlice["y"::-1], "yxdcb"), # type: ignore[misc]
- (pd.IndexSlice["y"::-4], "yb"), # type: ignore[misc]
- # absent labels
- (pd.IndexSlice[:"a":-1], "yxdcb"), # type: ignore[misc]
- (pd.IndexSlice[:"a":-2], "ydb"), # type: ignore[misc]
- (pd.IndexSlice["z"::-1], "yxdcb"), # type: ignore[misc]
- (pd.IndexSlice["z"::-3], "yc"), # type: ignore[misc]
- (pd.IndexSlice["m"::-1], "dcb"), # type: ignore[misc]
- (pd.IndexSlice[:"m":-1], "yx"), # type: ignore[misc]
- (pd.IndexSlice["a":"a":-1], ""), # type: ignore[misc]
- (pd.IndexSlice["z":"z":-1], ""), # type: ignore[misc]
- (pd.IndexSlice["m":"m":-1], ""), # type: ignore[misc]
- ],
- )
- def test_slice_locs_negative_step(self, in_slice, expected, any_string_dtype):
- index = Index(list("bcdxy"), dtype=any_string_dtype)
- s_start, s_stop = index.slice_locs(in_slice.start, in_slice.stop, in_slice.step)
- result = index[s_start : s_stop : in_slice.step]
- expected = Index(list(expected), dtype=any_string_dtype)
- tm.assert_index_equal(result, expected)
- def test_slice_locs_negative_step_oob(self, any_string_dtype):
- index = Index(list("bcdxy"), dtype=any_string_dtype)
- result = index[-10:5:1]
- tm.assert_index_equal(result, index)
- result = index[4:-10:-1]
- expected = Index(list("yxdcb"), dtype=any_string_dtype)
- tm.assert_index_equal(result, expected)
- def test_slice_locs_dup(self, any_string_dtype):
- index = Index(["a", "a", "b", "c", "d", "d"], dtype=any_string_dtype)
- assert index.slice_locs("a", "d") == (0, 6)
- assert index.slice_locs(end="d") == (0, 6)
- assert index.slice_locs("a", "c") == (0, 4)
- assert index.slice_locs("b", "d") == (2, 6)
- index2 = index[::-1]
- assert index2.slice_locs("d", "a") == (0, 6)
- assert index2.slice_locs(end="a") == (0, 6)
- assert index2.slice_locs("d", "b") == (0, 4)
- assert index2.slice_locs("c", "a") == (2, 6)
|