diff --git a/plotnine/_utils/__init__.py b/plotnine/_utils/__init__.py index cc83ee548..0aebd502b 100644 --- a/plotnine/_utils/__init__.py +++ b/plotnine/_utils/__init__.py @@ -355,8 +355,11 @@ def len_unique(x): def _id_var(x: AnyArrayLike, drop: bool = False) -> list[int]: """ - Assign ids to items in x. If two items - are the same, they get the same id. + Assign ids to items in x + + If two items are the same, they get the same id. + The ids start at 1 and, for categorical data, NaNs + get the highest id. Parameters ---------- @@ -364,25 +367,32 @@ def _id_var(x: AnyArrayLike, drop: bool = False) -> list[int]: items to associate ids with drop : bool Whether to drop unused factor levels + + Returns + ------- + ids: + List of ids """ if len(x) == 0: return [] if isinstance(x, pd.Series) and array_kind.categorical(x): + # The ids are a "re-coding" of the categorical codes/levels + # to meet the output requirements. if drop: x = x.cat.remove_unused_categories() - lst = list(x.cat.codes + 1) - else: - has_nan = any(np.isnan(i) for i in x if isinstance(i, float)) - if has_nan: - # NaNs are -1, we give them the highest code - nan_code = -1 - new_nan_code = np.max(x.cat.codes) + 1 - # TODO: We are assuming that x is of type Sequence[int|nan] - # is that accurate. - lst = [val if val != nan_code else new_nan_code for val in x] - else: - lst = list(x.cat.codes + 1) + + codes = x.cat.codes + + # We want our list to start at 1. + # But NaNs are -1, and if we have them, we want them to have + # the highest code, i.e. to be ordered last. + ids = list(codes + 1) + has_nan = (codes == -1).any() + if has_nan: + # The NaNs now have an id of 0 + highest_id = max(ids) + 1 + ids = [highest_id if i == 0 else i for i in ids] else: try: levels = sorted(set(x)) @@ -390,10 +400,9 @@ def _id_var(x: AnyArrayLike, drop: bool = False) -> list[int]: # x probably has NANs levels = multitype_sort(list(set(x))) - lst = match(x, levels) - lst = [item + 1 for item in lst] + ids = list(match(x, levels, start=1)) - return lst + return ids def join_keys(x, y, by=None): diff --git a/tests/test_facets.py b/tests/test_facets.py index 1410a5c38..167956d75 100644 --- a/tests/test_facets.py +++ b/tests/test_facets.py @@ -222,6 +222,25 @@ def test_facetting_with_unused_categories(): p.draw_test() # pyright: ignore +def test_facet_grid_with_missing_categorical_values(): + data = pd.DataFrame( + { + "x": [0.0] * 6, + "y": [0.0] * 6, + "row": ["a", "a", "a", "b", "b", "b"], + "col": pd.Categorical( + ["x", "y", None, "x", "y", None], + categories=["x", "y"], + ), + } + ) + + p = ggplot(data, aes("x", "y")) + geom_point() + facet_grid("row", "col") + + # No exception + p.draw_test() # pyright: ignore + + def test_invalid_scales_value_raises(): with pytest.raises(ValueError): # note the missing underscore diff --git a/tests/test_utils.py b/tests/test_utils.py index a7efdc09e..eeb569907 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -131,6 +131,32 @@ def test_ninteraction(): assert ninteraction(data) == [1] +def test_ninteraction_drops_unused_categorical_levels_with_missing(): + data = pd.DataFrame( + { + "a": pd.Categorical( + ["x", "y", None, "x"], + categories=["x", "unused", "y"], + ) + } + ) + + assert ninteraction(data, drop=True) == [1, 2, 3, 1] + + +def test_ninteraction_categorical_missing_values_get_highest_id(): + data = pd.DataFrame( + { + "a": pd.Categorical( + ["x", "y", None, "x"], + categories=["x", "unused", "y"], + ) + } + ) + + assert ninteraction(data, drop=False) == [1, 3, 4, 1] + + def test_ninteraction_datetime_series(): # When a pandas datetime is converted Numpy datetime, the two # no longer compare as equal! This test ensures that case is