示例#1
0
    def test_get_loc(self):
        # GH 12531
        cidx1 = CategoricalIndex(list("abcde"), categories=list("edabc"))
        idx1 = Index(list("abcde"))
        assert cidx1.get_loc("a") == idx1.get_loc("a")
        assert cidx1.get_loc("e") == idx1.get_loc("e")

        for i in [cidx1, idx1]:
            with pytest.raises(KeyError, match="'NOT-EXIST'"):
                i.get_loc("NOT-EXIST")

        # non-unique
        cidx2 = CategoricalIndex(list("aacded"), categories=list("edabc"))
        idx2 = Index(list("aacded"))

        # results in bool array
        res = cidx2.get_loc("d")
        tm.assert_numpy_array_equal(res, idx2.get_loc("d"))
        tm.assert_numpy_array_equal(
            res, np.array([False, False, False, True, False, True])
        )
        # unique element results in scalar
        res = cidx2.get_loc("e")
        assert res == idx2.get_loc("e")
        assert res == 4

        for i in [cidx2, idx2]:
            with pytest.raises(KeyError, match="'NOT-EXIST'"):
                i.get_loc("NOT-EXIST")

        # non-unique, sliceable
        cidx3 = CategoricalIndex(list("aabbb"), categories=list("abc"))
        idx3 = Index(list("aabbb"))

        # results in slice
        res = cidx3.get_loc("a")
        assert res == idx3.get_loc("a")
        assert res == slice(0, 2, None)

        res = cidx3.get_loc("b")
        assert res == idx3.get_loc("b")
        assert res == slice(2, 5, None)

        for i in [cidx3, idx3]:
            with pytest.raises(KeyError, match="'c'"):
                i.get_loc("c")
示例#2
0
 def test_get_loc_nonmonotonic_nonunique(self):
     cidx = CategoricalIndex(list("abcb"))
     result = cidx.get_loc("b")
     expected = np.array([False, True, False, True], dtype=bool)
     tm.assert_numpy_array_equal(result, expected)
示例#3
0
 def test_get_loc_monotonic_nonunique(self):
     cidx = CategoricalIndex(list("abbc"))
     result = cidx.get_loc("b")
     expected = slice(1, 3, None)
     assert result == expected
示例#4
0
 def test_get_loc_unique(self):
     cidx = CategoricalIndex(list("abc"))
     result = cidx.get_loc("b")
     assert result == 1
示例#5
0
    def test_get_loc_nan(self):
        # GH#41933
        ci = CategoricalIndex(["A", "B", np.nan])
        res = ci.get_loc(np.nan)

        assert res == 2