split test so that it actually render points with mpl

Author: timtreis

src/spatialdata_plot/pl/render.py

@@ -369,12 +369,20 @@ 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 (type(color_vector) is pd.core.arrays.categorical.Categorical)
-            and (len(color_vector.categories.values) > 1)
-            else None
-        )
+        color_key: dict[str, str] | None = None
+        if color_by_categorical and col_for_color is not None:
+            cat_series = pd.Categorical(transformed_element[col_for_color])
+            colors_arr = np.asarray(color_vector, dtype=object)
+            color_key = {}
+            for cat in cat_series.categories:
+                if cat == "nan":
+                    key_color = render_params.cmap_params.na_color.get_hex()
+                else:
+                    idx = np.flatnonzero(cat_series == cat)
+                    key_color = colors_arr[idx[0]] if idx.size else render_params.cmap_params.na_color.get_hex()
+                if isinstance(key_color, str) and key_color.startswith("#"):
+                    key_color = _hex_no_alpha(key_color)
+                color_key[str(cat)] = key_color
 
         if color_by_categorical or col_for_color is None:
             ds_cmap = None
@@ -897,16 +905,20 @@ 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 (type(color_vector) is pd.core.arrays.categorical.Categorical)
-            and (len(color_vector.categories.values) > 1)
-            else None
-        )
-
-        # remove alpha from color if it's hex
-        if color_key is not None and color_key[0][0] == "#":
-            color_key = [_hex_no_alpha(x) for x in color_key]
+        color_key: dict[str, str] | None = None
+        if color_by_categorical and col_for_color is not None:
+            cat_series = pd.Categorical(transformed_element[col_for_color])
+            colors_arr = np.asarray(color_vector, dtype=object)
+            color_key = {}
+            for cat in cat_series.categories:
+                if cat == "nan":
+                    key_color = render_params.cmap_params.na_color.get_hex()
+                else:
+                    idx = np.flatnonzero(cat_series == cat)
+                    key_color = colors_arr[idx[0]] if idx.size else render_params.cmap_params.na_color.get_hex()
+                if isinstance(key_color, str) and key_color.startswith("#"):
+                    key_color = _hex_no_alpha(key_color)
+                color_key[str(cat)] = key_color
         if isinstance(color_vector[0], str) and (color_vector is not None and color_vector[0][0] == "#"):
             color_vector = np.asarray([_hex_no_alpha(x) for x in color_vector])
 

src/spatialdata_plot/pl/utils.py

@@ -3046,7 +3046,7 @@ def _prepare_transformation(
 def _datashader_map_aggregate_to_color(
     agg: DataArray,
     cmap: str | list[str] | ListedColormap,
-    color_key: None | list[str] = None,
+    color_key: None | list[str] | Mapping[str, str] = None,
     min_alpha: float = 40,
     span: None | list[float] = None,
     clip: bool = True,

tests/pl/test_render_points.py

@@ -522,10 +522,12 @@ def test_plot_can_annotate_points_with_table_layer(self, sdata_blobs: SpatialDat
 
         sdata_blobs.pl.render_points("blobs_points", color="feature0", size=10, table_layer="normalized").pl.show()
 
-    def test_plot_can_annotate_points_with_nan_in_table_obs_categorical(
+    def test_plot_can_annotate_points_with_nan_in_table_obs_categorical_matplotlib(
         self, sdata_blobs_points_with_nans_in_table: SpatialData
     ):
-        sdata_blobs_points_with_nans_in_table.pl.render_points("blobs_points", color="category", size=30).pl.show()
+        sdata_blobs_points_with_nans_in_table.pl.render_points(
+            "blobs_points", color="category", size=40, method="matplotlib"
+        ).pl.show()
 
     def test_plot_can_annotate_points_with_nan_in_table_obs_categorical_datashader(
         self, sdata_blobs_points_with_nans_in_table: SpatialData