Fix CI (dask pinning, test with rasterize, mypy) (#809)

passing tests and pre-commits

Author: LucaMarconato

pyproject.toml

@@ -25,7 +25,7 @@ dependencies = [
     "anndata>=0.9.1",
     "click",
     "dask-image",
-    "dask>=2024.4.1",
+    "dask>=2024.4.1,<=2024.11.2",
     "fsspec",
     "geopandas>=0.14",
     "multiscale_spatial_image>=2.0.2",

src/spatialdata/_core/operations/rasterize.py

@@ -217,9 +217,9 @@ def rasterize(
         The table optionally containing the `value_key` and the name of the table in the returned `SpatialData` object.
         Must be `None` when `data` is a `SpatialData` object, otherwise it assumes the default value of `'table'`.
     return_regions_as_labels
-        By default, single-scale images of shape `(c, y, x)` are returned. If `True`, returns labels and shapes as
-        labels of shape `(y, x)` as opposed to an image of shape `(c, y, x)`. Points and images are always returned
-        as images, and multiscale raster data is always returned as single-scale data.
+        By default, single-scale images of shape `(c, y, x)` are returned. If `True`, returns labels, shapes and points
+        as labels of shape `(y, x)` as opposed to an image of shape `(c, y, x)`. Images are always returned as images,
+        and multiscale raster data is always returned as single-scale data.
     agg_func
         Available only when rasterizing points and shapes. A reduction function from datashader (its name, or a
         `Callable`). See the notes for more details on the default behavior.
@@ -234,6 +234,11 @@ def rasterize(
     into a `DataArray` (not a `DataTree`). So if a `SpatialData` object with elements is passed, a `SpatialData` object
     with single-scale images and labels will be returned.
 
+    When `return_regions_as_labels` is `True`, the returned `DataArray` object will have an attribute called
+    `label_index_to_category` that maps the label index to the category name. You can access it via
+    `returned_data.attrs["label_index_to_category"]`. The returned labels will start from 1 (0 is reserved for the
+    background), and will be contiguous.
+
     Notes
     -----
     For images and labels, the parameters `value_key`, `table_name`, `agg_func`, and `return_single_channel` are not
@@ -587,7 +592,7 @@ def rasterize_images_labels(
     )
     assert isinstance(transformed_dask, DaskArray)
     channels = xdata.coords["c"].values if schema in (Image2DModel, Image3DModel) else None
-    transformed_data = schema.parse(transformed_dask, dims=xdata.dims, c_coords=channels)  # type: ignore[call-arg,arg-type]
+    transformed_data = schema.parse(transformed_dask, dims=xdata.dims, c_coords=channels)  # type: ignore[call-arg]
 
     if target_coordinate_system != "global":
         remove_transformation(transformed_data, "global")
@@ -650,7 +655,7 @@ def rasterize_shapes_points(
     if value_key is not None:
         kwargs = {"sdata": sdata, "element_name": element_name} if element_name is not None else {"element": data}
         data[VALUES_COLUMN] = get_values(value_key, table_name=table_name, **kwargs).iloc[:, 0]  # type: ignore[arg-type, union-attr]
-    elif isinstance(data, GeoDataFrame):
+    elif isinstance(data, GeoDataFrame) or isinstance(data, DaskDataFrame) and return_regions_as_labels is True:
         value_key = VALUES_COLUMN
         data[VALUES_COLUMN] = data.index.astype("category")
     else:
@@ -706,6 +711,14 @@ def rasterize_shapes_points(
     agg = agg.fillna(0)
 
     if return_regions_as_labels:
+        if label_index_to_category is not None:
+            max_label = next(iter(reversed(label_index_to_category.keys())))
+        else:
+            max_label = int(agg.max().values)
+        max_uint16 = np.iinfo(np.uint16).max
+        if max_label > max_uint16:
+            raise ValueError(f"Maximum label index is {max_label}. Values higher than {max_uint16} are not supported.")
+        agg = agg.astype(np.uint16)
         return Labels2DModel.parse(agg, transformations=transformations)
 
     agg = agg.expand_dims(dim={"c": 1}).transpose("c", "y", "x")

src/spatialdata/_core/operations/transform.py

@@ -52,7 +52,7 @@ def _transform_raster(
     c_shape: tuple[int, ...]
     c_shape = (data.shape[0],) if "c" in axes else ()
     new_spatial_shape = tuple(
-        int(np.max(new_v[:, i]) - np.min(new_v[:, i])) for i in range(len(c_shape), n_spatial_dims + len(c_shape))  # type: ignore[operator]
+        int(np.max(new_v[:, i]) - np.min(new_v[:, i])) for i in range(len(c_shape), n_spatial_dims + len(c_shape))
     )
     output_shape = c_shape + new_spatial_shape
     translation_vector = np.min(new_v[:, :-1], axis=0)
@@ -86,8 +86,8 @@ def _transform_raster(
             # min_y_inverse = np.min(new_v_inverse[:, 1])
 
             if "c" in axes:
-                plt.imshow(da.moveaxis(transformed_dask, 0, 2), origin="lower", alpha=0.5)  # type: ignore[attr-defined]
-                plt.imshow(da.moveaxis(im, 0, 2), origin="lower", alpha=0.5)  # type: ignore[attr-defined]
+                plt.imshow(da.moveaxis(transformed_dask, 0, 2), origin="lower", alpha=0.5)
+                plt.imshow(da.moveaxis(im, 0, 2), origin="lower", alpha=0.5)
             else:
                 plt.imshow(transformed_dask, origin="lower", alpha=0.5)
                 plt.imshow(im, origin="lower", alpha=0.5)
@@ -322,7 +322,7 @@ def _(
     )
     c_coords = data.indexes["c"].values if "c" in data.indexes else None
     # mypy thinks that schema could be ShapesModel, PointsModel, ...
-    transformed_data = schema.parse(transformed_dask, dims=axes, c_coords=c_coords)  # type: ignore[call-arg,arg-type]
+    transformed_data = schema.parse(transformed_dask, dims=axes, c_coords=c_coords)  # type: ignore[call-arg]
     assert isinstance(transformed_data, DataArray)
     old_transformations = get_transformation(data, get_all=True)
     assert isinstance(old_transformations, dict)
@@ -448,7 +448,7 @@ def _(
     for ax in axes:
         indices = xtransformed["dim"] == ax
         new_ax = xtransformed[:, indices]
-        transformed[ax] = new_ax.data.flatten()  # type: ignore[attr-defined]
+        transformed[ax] = new_ax.data.flatten()
 
     old_transformations = get_transformation(data, get_all=True)
     assert isinstance(old_transformations, dict)
@@ -481,9 +481,9 @@ def _(
     )
     # TODO: nitpick, mypy expects a listof literals and here we have a list of strings.
     # I ignored but we may want to fix this
-    affine = transformation.to_affine(axes, axes)  # type: ignore[arg-type]
+    affine = transformation.to_affine(axes, axes)
     matrix = affine.matrix
-    shapely_notation = matrix[:-1, :-1].ravel().tolist() + matrix[:-1, -1].tolist()
+    shapely_notation = matrix[:-1, :-1].ravel().tolist() + matrix[:-1, -1].tolist()  # type: ignore[operator]
     transformed_geometry = data.geometry.affine_transform(shapely_notation)
     transformed_data = data.copy(deep=True)
     transformed_data.attrs[TRANSFORM_KEY] = {DEFAULT_COORDINATE_SYSTEM: Identity()}

src/spatialdata/_core/query/relational_query.py

@@ -214,7 +214,7 @@ def _filter_table_by_elements(
             # some instances have not a corresponding row in the table
             instances = np.setdiff1d(instances, n0)
         assert np.sum(to_keep) == len(instances)
-        assert sorted(set(instances.tolist())) == sorted(set(table.obs[instance_key].tolist()))
+        assert sorted(set(instances.tolist())) == sorted(set(table.obs[instance_key].tolist()))  # type: ignore[type-var]
         table_df = pd.DataFrame({instance_key: table.obs[instance_key], "position": np.arange(len(instances))})
         merged = pd.merge(table_df, pd.DataFrame(index=instances), left_on=instance_key, right_index=True, how="right")
         matched_positions = merged["position"].to_numpy()
@@ -467,7 +467,11 @@ def _left_join_spatialelement_table(
                 )
                 continue
 
-    joined_indices = joined_indices.dropna() if joined_indices is not None else None
+    if joined_indices is not None:
+        joined_indices = joined_indices.dropna()
+        # if nan were present, the dtype would have been changed to float
+        if joined_indices.dtype == float:
+            joined_indices = joined_indices.astype(int)
     joined_table = table[joined_indices, :].copy() if joined_indices is not None else None
     _inplace_fix_subset_categorical_obs(subset_adata=joined_table, original_adata=table)
 

src/spatialdata/_core/query/spatial_query.py

@@ -700,8 +700,8 @@ def _(
             bounding_box_mask = _bounding_box_mask_points(
                 points=points_query_coordinate_system,
                 axes=axes,
-                min_coordinate=min_c,
-                max_coordinate=max_c,
+                min_coordinate=min_c,  # type: ignore[arg-type]
+                max_coordinate=max_c,  # type: ignore[arg-type]
             )
             if len(bounding_box_mask) == 1:
                 bounding_box_mask = bounding_box_mask[0]

src/spatialdata/_io/io_points.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 
 import zarr
-from dask.dataframe import DataFrame as DaskDataFrame  # type: ignore[attr-defined]
+from dask.dataframe import DataFrame as DaskDataFrame
 from dask.dataframe import read_parquet
 from ome_zarr.format import Format
 

src/spatialdata/_types.py

@@ -1,18 +1,14 @@
+from typing import Any
+
 import numpy as np
 from xarray import DataArray, DataTree
 
 __all__ = ["ArrayLike", "ColorLike", "DTypeLike", "Raster_T"]
 
-try:
-    from numpy.typing import DTypeLike, NDArray
-
-    ArrayLike = NDArray[np.float64]
-    IntArrayLike = NDArray[np.int64]  # or any np.integer
+from numpy.typing import DTypeLike, NDArray
 
-except (ImportError, TypeError):
-    ArrayLike = np.ndarray  # type: ignore[misc]
-    IntArrayLike = np.ndarray  # type: ignore[misc]
-    DTypeLike = np.dtype  # type: ignore[misc, assignment]
+ArrayLike = NDArray[np.floating[Any]]
+IntArrayLike = NDArray[np.integer[Any]]
 
 Raster_T = DataArray | DataTree
 ColorLike = tuple[float, ...] | str

src/spatialdata/_utils.py

@@ -80,7 +80,7 @@ def _compute_paddings(data: DataArray, axis: str) -> tuple[int, int]:
         others = list(data.dims)
         others.remove(axis)
         # mypy (luca's pycharm config) can't see the isclose method of dask array
-        s = da.isclose(data.sum(dim=others), 0)  # type: ignore[attr-defined]
+        s = da.isclose(data.sum(dim=others), 0)
         # TODO: rewrite this to use dask array; can't get it to work with it
         x = s.compute()
         non_zero = np.where(x == 0)[0]

src/spatialdata/dataloader/datasets.py

@@ -144,7 +144,7 @@ def __init__(
                 **dict(rasterize_kwargs),
             )
             if rasterize
-            else bounding_box_query  # type: ignore[assignment]
+            else bounding_box_query
         )
         self._return = self._get_return(return_annotations, table_name)
         self.transform = transform

src/spatialdata/datasets.py

@@ -182,7 +182,7 @@ def _image_blobs(
         masks = []
         for i in range(n_channels):
             mask = self._generate_blobs(length=length, seed=i)
-            mask = (mask - mask.min()) / np.ptp(mask)  # type: ignore[attr-defined]
+            mask = (mask - mask.min()) / np.ptp(mask)
             masks.append(mask)
 
         x = np.stack(masks, axis=0)