fix(datashader): route none-state shapes through the categorical color-key (#700)

Review caught a regression from the earlier is_categorical switch: shapes
datashader keyed color_by_categorical on is_categorical, dropping the none
state (all-NaN color column) into the numeric reduction → ValueError: axes
don't match array. Restore the old truth set (categorical OR none when a color
column is set) via `col_for_color is not None and not is_continuous`; add a
datashader regression test. (Points already kept the none state.)

Cleanups from the same review:
- labels as_points: classify the fresh point-color spec by its own source
  vector instead of reusing the parent colortype (was constructing
  invariant-violating none/categorical specs with source=None).
- _get_collection_shape: pass the already-(N,4) RGBA fill through instead of a
  redundant to_rgba_array round-trip (dropped a full-N copy on the hot path).
- to_rgba: compute na_rgba only on the numeric branch that uses it.
- collapse with_color_vector/with_source_vector into one evolve(**changes).

Author: timtreis

src/spatialdata_plot/pl/_color.py

@@ -702,13 +702,9 @@ def apply_transfunc(self, transfunc: Any) -> ColorSpec:
             return replace(self, color_vector=transfunc(self.color_vector))
         return self
 
-    def with_color_vector(self, color_vector: ArrayLike) -> ColorSpec:
-        """Return a copy with a replaced ``color_vector`` (a post-resolution rewrite: reprocess, compute)."""
-        return replace(self, color_vector=color_vector)
-
-    def with_source_vector(self, source_vector: ArrayLike | pd.Series | None) -> ColorSpec:
-        """Return a copy with a replaced ``source_vector`` (e.g. ``remove_unused_categories``, ``compute``)."""
-        return replace(self, source_vector=source_vector)
+    def evolve(self, **changes: Any) -> ColorSpec:
+        """Return a copy with ``source_vector``/``color_vector`` replaced (post-resolution rewrites)."""
+        return replace(self, **changes)
 
     def make_palette(self) -> ListedColormap:
         """Build a ``ListedColormap`` from the colors, dropping NaN categories when categorical."""
@@ -734,14 +730,13 @@ def to_rgba(self, cmap_params: CmapParams) -> np.ndarray:
         Categorical/none ``color_vector`` is per-row hex -> straight to RGBA; continuous numerics map
         via norm+cmap with NaN/non-finite rows painted ``na_color``; an object vector mixes the two.
         """
-        na_rgba = colors.to_rgba(cmap_params.na_color.get_hex_with_alpha())
         if self.source_vector is not None:  # categorical or none: color_vector holds per-row hex
             return np.asarray(ColorConverter().to_rgba_array(list(self.color_vector)))
         arr = np.asarray(self.color_vector)
         if arr.ndim == 2 and arr.shape[1] in (3, 4) and np.issubdtype(arr.dtype, np.number):
             return np.asarray(ColorConverter().to_rgba_array(arr))
         rgba = np.empty((len(arr), 4), dtype=float)
-        rgba[:] = na_rgba
+        rgba[:] = colors.to_rgba(cmap_params.na_color.get_hex_with_alpha())
         if np.issubdtype(arr.dtype, np.number):
             finite = np.isfinite(arr)
             if finite.any():

src/spatialdata_plot/pl/_geometry.py

@@ -180,16 +180,12 @@ def _get_collection_shape(
     color / list of color specs (broadcast). ``outline_color`` may be an ``(N, 4)`` float RGBA array,
     whose alpha channel is mutated in place to apply ``outline_alpha`` (pass a copy to retain it).
     """
-    # Per-row fill is pre-mapped to RGBA by ``ColorSpec.to_rgba``; here ``c`` is either that
-    # (N, 4) array or a single color / list of color specs (e.g. the white outline placeholder).
+    # Per-row fill is pre-mapped to RGBA by ``ColorSpec.to_rgba`` (an (N, 4) float array, used
+    # as-is); otherwise ``c`` is a single color / list of color specs (e.g. the white outline
+    # placeholder), which ``to_rgba_array`` broadcasts.
     c_arr = np.asarray(c)
-    if (
-        c_arr.ndim == 2
-        and c_arr.shape[0] == len(shapes)
-        and c_arr.shape[1] in (3, 4)
-        and np.issubdtype(c_arr.dtype, np.number)
-    ):
-        fill_c = np.asarray(ColorConverter().to_rgba_array(c_arr))
+    if c_arr.shape == (len(shapes), 4) and np.issubdtype(c_arr.dtype, np.floating):
+        fill_c = c_arr
     else:
         fill_c = np.asarray(ColorConverter().to_rgba_array(c))
 

src/spatialdata_plot/pl/render.py

@@ -684,7 +684,7 @@ def _render_shapes(
     norm = render_params.cmap_params.fresh_norm()
 
     if len(color_spec.color_vector) == 0:
-        color_spec = color_spec.with_color_vector([render_params.cmap_params.na_color.get_hex_with_alpha()])
+        color_spec = color_spec.evolve(color_vector=[render_params.cmap_params.na_color.get_hex_with_alpha()])
 
     # continuous case: leave NaNs as NaNs; utils maps them to na_color during draw
     if color_spec.is_continuous:
@@ -708,7 +708,7 @@ def _render_shapes(
                     f"Found {nan_count} NaN values in color data. "
                     "These observations will be colored with the 'na_color'."
                 )
-        color_spec = color_spec.with_color_vector(cv)
+        color_spec = color_spec.evolve(color_vector=cv)
 
     palette = color_spec.make_palette()
 
@@ -718,7 +718,7 @@ def _render_shapes(
     )
     if has_valid_color and color_spec.is_categorical and col_for_color is not None:
         # necessary in case different shapes elements are annotated with one table
-        color_spec = color_spec.with_source_vector(color_spec.source_vector.remove_unused_categories())
+        color_spec = color_spec.evolve(source_vector=color_spec.source_vector.remove_unused_categories())
 
     shapes = gpd.GeoDataFrame(shapes, geometry="geometry")
 
@@ -816,7 +816,7 @@ def _render_shapes(
 
         cvs = ds.Canvas(plot_width=plot_width, plot_height=plot_height, x_range=x_ext, y_range=y_ext)
 
-        # datashader consumes raw arrays; unpack the carrier for this leaf and re-wrap its result
+        # datashader leaf consumes raw arrays; re-wrapped into the carrier after aggregation
         color_vector = color_spec.color_vector
         color_source_vector = color_spec.source_vector
 
@@ -839,8 +839,9 @@ def _render_shapes(
                         color_vector = list(color_vector) + [na_color] * (len(transformed_element) - len(color_vector))
 
             transformed_element[col_for_color] = color_vector if color_source_vector is None else color_source_vector
-        # Render shapes with datashader
-        color_by_categorical = color_spec.is_categorical
+        # Render shapes with datashader: categorical AND none (na-array) go through the categorical
+        # color-key path; only continuous uses the numeric reduction.
+        color_by_categorical = col_for_color is not None and not color_spec.is_continuous
         if color_by_categorical:
             cat_series = transformed_element[col_for_color]
             if not isinstance(cat_series.dtype, pd.CategoricalDtype):
@@ -861,7 +862,7 @@ def _render_shapes(
             default_reduction=_default_reduction,
             kind="shapes",
         )
-        color_spec = color_spec.with_color_vector(color_vector)
+        color_spec = color_spec.evolve(color_vector=color_vector)
 
         _render_ds_outlines(
             cvs,
@@ -1106,7 +1107,7 @@ def _render_centroids_as_points(
             as_markers=True,
             axes_extent=axes_extent,
         )
-        color_spec = color_spec.with_source_vector(csv).with_color_vector(cv)
+        color_spec = color_spec.evolve(source_vector=csv, color_vector=cv)
     else:
         cax = _scatter_points(
             ax,
@@ -1465,7 +1466,7 @@ def _render_points(
             fig_params=fig_params,
             default_reduction=_default_reduction,
         )
-        color_spec = color_spec.with_source_vector(csv).with_color_vector(cv)
+        color_spec = color_spec.evolve(source_vector=csv, color_vector=cv)
 
     elif method == "matplotlib":
         # update axis limits if plot was empty before (necessary if datashader comes after)
@@ -2319,7 +2320,12 @@ def _render_labels(
             render_params,
             x=xy[:, 0],
             y=xy[:, 1],
-            color_spec=ColorSpec(color_spec.colortype, point_color_source_vector, point_color_vector),
+            # point colours are derived fresh; classify by source so the spec stays self-consistent
+            color_spec=ColorSpec(
+                "categorical" if point_color_source_vector is not None else "continuous",
+                point_color_source_vector,
+                point_color_vector,
+            ),
             norm=render_params.cmap_params.fresh_norm(),  # ax.scatter autoscales in place; don't mutate the shared norm
             na_color=na_color,
             adata=table if table_name is not None else None,

tests/pl/test_render_shapes.py

@@ -1132,6 +1132,15 @@ def test_render_shapes_all_nan_outline_color_does_not_crash(sdata_blobs_shapes_a
     plt.close(fig)
 
 
+def test_render_shapes_all_nan_color_datashader_does_not_crash(sdata_blobs_shapes_annotated: SpatialData):
+    # Regression: all-NaN color is the "none" colortype; the datashader path must route it through
+    # the categorical color-key (na_color), not the numeric reduction (which rejects the na column).
+    sdata_blobs_shapes_annotated["blobs_polygons"]["nanvals"] = [np.nan] * 5
+    fig, ax = plt.subplots()
+    sdata_blobs_shapes_annotated.pl.render_shapes("blobs_polygons", color="nanvals", method="datashader").pl.show(ax=ax)
+    plt.close(fig)
+
+
 def test_render_shapes_all_nan_color_with_groups_does_not_crash(sdata_blobs_shapes_annotated: SpatialData):
     # Regression: all-NaN color is the "none" colortype (na-array source, not categorical); the
     # groups preamble must not feed it to _warn_missing_groups, which reads `.categories`.