Fix datashader failing with single-category Categorical (#539)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: timtreis

src/spatialdata_plot/pl/render.py

@@ -62,6 +62,21 @@
 _Normalize = Normalize | abc.Sequence[Normalize]
 
 
+def _build_color_key_from_categorical(color_vector: pd.Categorical | np.ndarray | object) -> list[str] | None:
+    """Build a datashader ``color_key`` list from a categorical color vector.
+
+    Returns ``None`` when *color_vector* is not a :class:`pd.Categorical` or
+    has no categories.  Hex colours are stripped of their alpha channel;
+    named colours (e.g. ``"red"``) are passed through unchanged.
+    """
+    if not isinstance(getattr(color_vector, "dtype", None), pd.CategoricalDtype):
+        return None
+    cat_values = color_vector.categories.values  # type: ignore[union-attr]
+    if len(cat_values) == 0:
+        return None
+    return [_hex_no_alpha(x) if isinstance(x, str) and x.startswith("#") else x for x in cat_values]
+
+
 def _split_colorbar_params(params: dict[str, object] | None) -> tuple[dict[str, object], dict[str, object], str | None]:
     """Split colorbar params into layout hints, Matplotlib kwargs, and label override."""
     layout: dict[str, object] = {}
@@ -301,7 +316,7 @@ def _render_shapes(
         # Render shapes with datashader
         color_by_categorical = col_for_color is not None and color_source_vector is not None
         aggregate_with_reduction = None
-        if col_for_color is not None and (render_params.groups is None or len(render_params.groups) > 1):
+        if col_for_color is not None:
             if color_by_categorical:
                 agg = cvs.polygons(
                     transformed_element,
@@ -356,11 +371,7 @@ def _render_shapes(
                     agg = agg.where((agg <= norm.vmin) | (np.isnan(agg)), other=2)
                     agg = agg.where((agg != norm.vmin) | (np.isnan(agg)), other=0.5)
 
-        color_key = (
-            [_hex_no_alpha(x) for x in color_vector.categories.values]
-            if isinstance(color_vector.dtype, pd.CategoricalDtype) and (len(color_vector.categories.values) > 1)
-            else None
-        )
+        color_key = _build_color_key_from_categorical(color_vector)
 
         if color_by_categorical or col_for_color is None:
             ds_cmap = None
@@ -814,7 +825,7 @@ def _render_points(
         if color_by_categorical and transformed_element[col_for_color].values.dtype == object:
             transformed_element[col_for_color] = transformed_element[col_for_color].astype("category")
         aggregate_with_reduction = None
-        if col_for_color is not None and (render_params.groups is None or len(render_params.groups) > 1):
+        if col_for_color is not None:
             if color_by_categorical:
                 agg = cvs.points(transformed_element, "x", "y", agg=ds.by(col_for_color, ds.count()))
             else:
@@ -851,15 +862,7 @@ def _render_points(
                     agg = agg.where((agg <= norm.vmin) | (np.isnan(agg)), other=2)
                     agg = agg.where((agg != norm.vmin) | (np.isnan(agg)), other=0.5)
 
-        color_key: list[str] | None = (
-            list(color_vector.categories.values)
-            if isinstance(color_vector.dtype, pd.CategoricalDtype) and (len(color_vector.categories.values) > 1)
-            else None
-        )
-
-        # remove alpha from color if it's hex
-        if color_key is not None and all(len(x) == 9 for x in color_key) and color_key[0][0] == "#":
-            color_key = [x[:-2] for x in color_key]
+        color_key = _build_color_key_from_categorical(color_vector)
         if isinstance(color_vector[0], str) and (
             color_vector is not None and all(len(x) == 9 for x in color_vector) and color_vector[0][0] == "#"
         ):

tests/pl/test_render_points.py

@@ -572,3 +572,37 @@ def test_datashader_colors_points_from_table_obs(sdata_blobs: SpatialData):
         method="datashader",
         size=5,
     ).pl.show()
+
+
+def test_plot_datashader_single_category_points(sdata_blobs: SpatialData):
+    """Datashader with a single-category Categorical must not raise.
+
+    Regression test for https://github.com/scverse/spatialdata-plot/issues/483.
+    Before the fix, color_key was None when there was only 1 category, but ds.by()
+    still produced a 3D DataArray, causing datashader to raise:
+        ValueError: Color key must be provided, with at least as many colors as
+        there are categorical fields
+    """
+    n_obs = len(sdata_blobs["blobs_points"])
+    obs = pd.DataFrame(
+        {
+            "instance_id": np.arange(n_obs),
+            "region": pd.Categorical(["blobs_points"] * n_obs),
+            "foo": pd.Categorical(["only_cat"] * n_obs),
+        }
+    )
+    table = TableModel.parse(
+        adata=AnnData(get_standard_RNG().normal(size=(n_obs, 3)), obs=obs),
+        region="blobs_points",
+        region_key="region",
+        instance_key="instance_id",
+    )
+    sdata_blobs["single_cat_table"] = table
+
+    sdata_blobs.pl.render_points(
+        "blobs_points",
+        color="foo",
+        table_name="single_cat_table",
+        method="datashader",
+        size=5,
+    ).pl.show()

tests/pl/test_render_shapes.py

@@ -973,3 +973,29 @@ def test_plot_can_handle_mixed_numeric_and_color_data(sdata_blobs: SpatialData):
     # Mixed numeric / non-numeric values should raise a TypeError
     with pytest.raises(TypeError, match="contains both numeric and non-numeric values"):
         sdata_blobs.pl.render_shapes(element="blobs_circles", color="mixed_data", na_color="gray").pl.show()
+
+
+def test_plot_datashader_single_category(sdata_blobs: SpatialData):
+    """Datashader with a single-category Categorical must not raise.
+
+    Regression test for https://github.com/scverse/spatialdata-plot/issues/483.
+    Before the fix, color_key was None when there was only 1 category, but ds.by()
+    still produced a 3D DataArray, causing datashader to raise:
+        ValueError: Color key must be provided, with at least as many colors as
+        there are categorical fields
+    """
+    n_obs = len(sdata_blobs["blobs_polygons"])
+    adata = AnnData(get_standard_RNG().normal(size=(n_obs, 10)))
+    adata.obs = pd.DataFrame(get_standard_RNG().normal(size=(n_obs, 3)), columns=["a", "b", "c"])
+    adata.obs["category"] = pd.Categorical(["only_cat"] * n_obs)
+    adata.obs["instance_id"] = list(range(n_obs))
+    adata.obs["region"] = "blobs_polygons"
+    table = TableModel.parse(
+        adata=adata,
+        region_key="region",
+        instance_key="instance_id",
+        region="blobs_polygons",
+    )
+    sdata_blobs["table"] = table
+
+    sdata_blobs.pl.render_shapes(element="blobs_polygons", color="category", method="datashader").pl.show()