Fixes basic typing issues in rasterize.py

AGENTS.md

@@ -0,0 +1,67 @@
+# AGENTS.md
+
+## Typing guidelines
+
+- Avoid `typing.cast` and `# type: ignore`, especially as workarounds for narrowing problems; use the style above
+  instead. If one is truly unavoidable, keep it as narrow as possible and add a comment explaining why
+  (e.g. `# FIXME: ...`).
+- Avoid creating functions that return unions (except if it's something like `T | None`). If a function has different
+  paths to return different types, it should probably be two functions or at least an overload.
+
+### Union narrowing style
+
+Narrow unions by asserting what a value is **not**, then use `typing.assert_type` to state (and have the type checker
+verify) what it is:
+
+```python
+from typing import assert_type
+
+transformation: BaseTransformation | dict = ...
+
+# Don't
+assert isinstance(transformation, BaseTransformation)
+
+# Do
+assert not isinstance(transformation, dict)
+assert_type(transformation, BaseTransformation)
+```
+
+When all branches of a union have been handled, use `typing.assert_never` instead of a generic "should not reach here"
+error, so the type checker tells us if a new type is added to the union and not handled:
+
+```python
+from typing import assert_never
+
+if isinstance(data, DataArray):
+    ...
+elif isinstance(data, DataTree):
+    ...
+else:
+    assert_never(data)  # instead of: raise RuntimeError("Should not reach here")
+```
+
+#### Why
+
+Because of Python's (multiple-)inheritance model, a positive `isinstance` assertion can never be rejected by a type
+checker, even when it contradicts the declared type:
+
+```python
+a: int = 123
+assert isinstance(a, list)  # accepted; the checker infers `a: <subclass of int and list>`
+```
+
+So a positive assertion can silently force a value to a type it should never have.
+
+#### Examples
+
+```python
+transformations: BaseTransformation | dict[str, BaseTransformation] = get_transformation(element, get_all=True)
+assert not isinstance(transformations, BaseTransformation)
+assert_type(transformations, dict[str, BaseTransformation])
+```
+
+Optional values are already narrowed "negatively", so keep using `is not None`:
+
+```python
+assert target_depth is not None
+```

src/spatialdata/_core/operations/rasterize.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, NamedTuple, assert_never, assert_type, cast
 
 import numpy as np
 from dask.array import Array as DaskArray
 from dask.dataframe import DataFrame as DaskDataFrame
 from geopandas import GeoDataFrame
+from numpy.typing import DTypeLike
 from shapely import Point
 from xarray import DataArray, DataTree
 
@@ -46,6 +47,12 @@
 VALUES_COLUMN = "__values_column"
 
 
+class TargetDimensions(NamedTuple):
+    target_width: float
+    target_height: float
+    target_depth: float | None
+
+
 def _compute_target_dimensions(
     spatial_axes: tuple[str, ...],
     min_coordinate: ListOrNDArrayFloating,
@@ -54,7 +61,7 @@ def _compute_target_dimensions(
     target_width: float | None,
     target_height: float | None,
     target_depth: float | None,
-) -> tuple[float, float, float | None]:
+) -> TargetDimensions:
     """
     Compute the pixel sizes (width, height, depth) of the image that will be produced by the rasterization.
 
@@ -103,24 +110,22 @@ def _compute_target_dimensions(
         == 1
     ), "you must specify only one of: target_unit_to_pixels, target_width, target_height, target_depth"
     assert set(spatial_axes) == {"x", "y"} or set(spatial_axes) == {"x", "y", "z"}
-    if "z" not in spatial_axes:
-        assert target_depth is None, "you cannot specify a target depth for 2D data"
 
     x_index = spatial_axes.index("x")
     y_index = spatial_axes.index("y")
     w_bb = max_coordinate[x_index] - min_coordinate[x_index]
     h_bb = max_coordinate[y_index] - min_coordinate[y_index]
     assert w_bb > 0, "all max_coordinate values must be greater than all min_coordinate values"
     assert h_bb > 0, "all max_coordinate values must be greater than all min_coordinate values"
-    w_to_h_bb = w_bb / h_bb
+    w_to_h_bb = float(w_bb / h_bb)
 
     d_bb = None
     d_to_h_bb = None
     if "z" in spatial_axes:
         z_index = spatial_axes.index("z")
         d_bb = max_coordinate[z_index] - min_coordinate[z_index]
         assert d_bb > 0, "all max_coordinate values must be greater than all min_coordinate values"
-        d_to_h_bb = d_bb / h_bb
+        d_to_h_bb = float(d_bb / h_bb)
 
     if target_unit_to_pixels is not None:
         target_width = w_bb * target_unit_to_pixels
@@ -140,15 +145,18 @@ def _compute_target_dimensions(
             target_depth = target_height * d_to_h_bb
     elif target_depth is not None:
         assert d_to_h_bb is not None
+        assert "z" in spatial_axes, "you cannot specify a target depth for 2D data"
         target_height = target_depth / d_to_h_bb
         target_width = target_height * w_to_h_bb
     else:
         raise RuntimeError("Should not reach here")
-    assert target_width is not None
     assert isinstance(target_width, float)
-    assert target_height is not None
     assert isinstance(target_height, float)
-    return np.round(target_width), np.round(target_height), np.round(target_depth) if target_depth is not None else None
+    return TargetDimensions(
+        target_width=np.round(target_width),
+        target_height=np.round(target_height),
+        target_depth=np.round(target_depth) if target_depth is not None else None,
+    )
 
 
 def rasterize(
@@ -305,7 +313,8 @@ def rasterize(
                     return_single_channel=return_single_channel if element_type in ("points", "shapes") else None,
                 )
                 new_name = f"{name}_rasterized_{element_type}"
-                assert isinstance(rasterized, DataArray)
+                assert not isinstance(rasterized, SpatialData)
+                assert_type(rasterized, DataArray)
                 model = get_model(rasterized)
                 if model in (Image2DModel, Image3DModel):
                     new_images[new_name] = rasterized
@@ -335,7 +344,8 @@ def rasterize(
             target_depth=target_depth,
         )
         transformations = get_transformation(rasterized, get_all=True)
-        assert isinstance(transformations, dict)
+        assert not isinstance(transformations, BaseTransformation)
+        assert_type(transformations, dict[str, BaseTransformation])
         # adjust the return type
         if model in (Labels2DModel, Labels3DModel) and not return_regions_as_labels:
             model = Image2DModel if model == Labels2DModel else Image3DModel
@@ -346,14 +356,14 @@ def rasterize(
             kwargs = {"sdata": sdata, "element_name": element_name} if element_name is not None else {"element": data}
             values = get_values(value_key, table_name=table_name, **kwargs).iloc[:, 0]  # type: ignore[arg-type, union-attr]
             max_index: int = np.max(values.index)
-            assigner = np.zeros(max_index + 1, dtype=values.dtype)
+            dtype = cast(DTypeLike, values.dtype)  # FIXME: try to assert instead
+            assigner = np.zeros(max_index + 1, dtype=dtype)
             assigner[values.index] = values
-            # call-arg is ignored because model is never TableModel (the error is that the transformation param is not
-            # accepted by TableModel.parse)
-            rasterized = model.parse(assigner[rasterized], transformations=transformations)  # type: ignore[call-arg]
+            rasterized = model.parse(assigner[rasterized], transformations=transformations)
         return rasterized
     if model in (PointsModel, ShapesModel):
-        assert isinstance(parsed_data, GeoDataFrame | DaskDataFrame)
+        assert not isinstance(parsed_data, DataArray | DataTree)
+        assert_type(parsed_data, GeoDataFrame | DaskDataFrame)
         return rasterize_shapes_points(
             data=parsed_data,
             axes=axes,
@@ -414,12 +424,14 @@ def _get_xarray_data_to_rasterize(
         latest_scale: str | None = None
         for scale in reversed(list(data.keys())):
             data_tree = data[scale]
-            assert isinstance(data_tree, DataTree)
+            assert not isinstance(data_tree, DataArray)
+            assert_type(data_tree, DataTree)
             latest_scale = scale
             v = list(data_tree.values())
             assert len(v) == 1
             scale_variable = v[0]
-            assert isinstance(scale_variable, DataArray)
+            assert not isinstance(scale_variable, DataTree)
+            assert_type(scale_variable, DataArray)
             xdata = scale_variable
             assert set(get_spatial_axes(tuple(str(dim) for dim in xdata.sizes))) == set(axes)
 
@@ -457,17 +469,19 @@ def _get_xarray_data_to_rasterize(
                 break
         assert latest_scale is not None
         latest_scale_tree = data[latest_scale]
-        assert isinstance(latest_scale_tree, DataTree)
+        assert not isinstance(latest_scale_tree, DataArray)
+        assert_type(latest_scale_tree, DataTree)
         latest_variable = next(iter(latest_scale_tree.values()))
-        assert isinstance(latest_variable, DataArray)
+        assert not isinstance(latest_variable, DataTree)
+        assert_type(latest_variable, DataArray)
         xdata = latest_variable
         if latest_scale != "scale0":
             transformations = xdata.attrs["transform"]
             pyramid_scale = _get_scale(transformations)
         else:
             pyramid_scale = None
     else:
-        raise RuntimeError("Should not reach here")
+        assert_never(data)
     return xdata, pyramid_scale
 
 
@@ -491,7 +505,8 @@ def _get_corrected_affine_matrix(
 
     """
     transformation = get_transformation(data, target_coordinate_system)
-    assert isinstance(transformation, BaseTransformation)
+    assert not isinstance(transformation, dict)
+    assert_type(transformation, BaseTransformation)
     affine = _get_affine_for_element(data, transformation)
     target_axes_unordered = affine.output_axes
     assert set(target_axes_unordered) in [{"x", "y", "z"}, {"x", "y"}, {"c", "x", "y", "z"}, {"c", "x", "y"}]
@@ -548,6 +563,8 @@ def rasterize_images_labels(
     offset = [min_coordinate[axes.index(ax)] for ax in axes]
     translation = Translation(offset, axes=axes)
 
+    assert not isinstance(data, GeoDataFrame | DaskDataFrame)
+    assert_type(data, DataArray | DataTree)
     xdata, pyramid_scale = _get_xarray_data_to_rasterize(
         data=data,
         axes=axes,
@@ -620,7 +637,8 @@ def rasterize_images_labels(
     set_transformation(transformed_data, sequence, target_coordinate_system)
 
     computed = compute_coordinates(transformed_data)
-    assert isinstance(computed, DataArray)
+    assert not isinstance(computed, DataTree)
+    assert_type(computed, DataArray)
     schema.validate(computed)
     return computed