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)