def test_where(self, klass):
i = CategoricalIndex(list("aabbca"), categories=list("cab"), ordered=False)
cond = [True] * len(i)
expected = i
result = i.where(klass(cond))
tm.assert_index_equal(result, expected)
cond = [False] + [True] * (len(i) - 1)
expected = CategoricalIndex([np.nan] + i[1:].tolist(), categories=i.categories)
result = i.where(klass(cond))
tm.assert_index_equal(result, expected)
def test_where_non_categories(self):
ci = CategoricalIndex(["a", "b", "c", "d"])
mask = np.array([True, False, True, False])
msg = "Cannot setitem on a Categorical with a new category"
with pytest.raises(ValueError, match=msg):
ci.where(mask, 2)
with pytest.raises(ValueError, match=msg):
# Test the Categorical method directly
ci._data.where(mask, 2)
def test_where_non_categories(self):
ci = CategoricalIndex(["a", "b", "c", "d"])
mask = np.array([True, False, True, False])
result = ci.where(mask, 2)
expected = Index(["a", 2, "c", 2], dtype=object)
tm.assert_index_equal(result, expected)
msg = "Cannot setitem on a Categorical with a new category"
with pytest.raises(ValueError, match=msg):
# Test the Categorical method directly
ci._data.where(mask, 2)
