Skip to content

Commit f35bcac

Browse files
authored
TST: fix test creating invalid CategoricalBlock (#32519)
* TST: fix test creating invalid CategoricalBlock * ValueError->AssertionError
1 parent 3b66021 commit f35bcac

File tree

2 files changed

+12
-5
lines changed

2 files changed

+12
-5
lines changed

pandas/core/internals/blocks.py

+4
Original file line numberDiff line numberDiff line change
@@ -1739,6 +1739,10 @@ def __init__(self, values, placement, ndim=None):
17391739
values = self._maybe_coerce_values(values)
17401740
super().__init__(values, placement, ndim)
17411741

1742+
if self.ndim == 2 and len(self.mgr_locs) != 1:
1743+
# TODO(2DEA): check unnecessary with 2D EAs
1744+
raise AssertionError("block.size != values.size")
1745+
17421746
def _maybe_coerce_values(self, values):
17431747
"""
17441748
Unbox to an extension array.

pandas/tests/internals/test_internals.py

+8-5
Original file line numberDiff line numberDiff line change
@@ -540,6 +540,13 @@ def _compare(old_mgr, new_mgr):
540540
assert new_mgr.get("g").dtype == np.float64
541541
assert new_mgr.get("h").dtype == np.float16
542542

543+
def test_invalid_ea_block(self):
544+
with pytest.raises(AssertionError, match="block.size != values.size"):
545+
create_mgr("a: category; b: category")
546+
547+
with pytest.raises(AssertionError, match="block.size != values.size"):
548+
create_mgr("a: category2; b: category2")
549+
543550
def test_interleave(self):
544551

545552
# self
@@ -552,14 +559,10 @@ def test_interleave(self):
552559
# will be converted according the actual dtype of the underlying
553560
mgr = create_mgr("a: category")
554561
assert mgr.as_array().dtype == "i8"
555-
mgr = create_mgr("a: category; b: category")
556-
assert mgr.as_array().dtype == "i8"
557562
mgr = create_mgr("a: category; b: category2")
558563
assert mgr.as_array().dtype == "object"
559564
mgr = create_mgr("a: category2")
560565
assert mgr.as_array().dtype == "object"
561-
mgr = create_mgr("a: category2; b: category2")
562-
assert mgr.as_array().dtype == "object"
563566

564567
# combinations
565568
mgr = create_mgr("a: f8")
@@ -702,7 +705,7 @@ def test_equals(self):
702705
"a:i8;b:f8", # basic case
703706
"a:i8;b:f8;c:c8;d:b", # many types
704707
"a:i8;e:dt;f:td;g:string", # more types
705-
"a:i8;b:category;c:category2;d:category2", # categories
708+
"a:i8;b:category;c:category2", # categories
706709
"c:sparse;d:sparse_na;b:f8", # sparse
707710
],
708711
)

0 commit comments

Comments
 (0)