Replacing fixed size tuples as return types (#1043)

* change return types and update tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix mypy problems

* fix typing

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: selmanozleyen

src/squidpy/gr/__init__.py

@@ -2,10 +2,31 @@
 
 from __future__ import annotations
 
-from squidpy.gr._build import mask_graph, spatial_neighbors
+from squidpy.gr._build import SpatialNeighborsResult, mask_graph, spatial_neighbors
 from squidpy.gr._ligrec import ligrec
-from squidpy.gr._nhood import centrality_scores, interaction_matrix, nhood_enrichment
+from squidpy.gr._nhood import (
+    NhoodEnrichmentResult,
+    centrality_scores,
+    interaction_matrix,
+    nhood_enrichment,
+)
 from squidpy.gr._niche import calculate_niche
 from squidpy.gr._ppatterns import co_occurrence, spatial_autocorr
 from squidpy.gr._ripley import ripley
 from squidpy.gr._sepal import sepal
+
+__all__ = [
+    "SpatialNeighborsResult",
+    "NhoodEnrichmentResult",
+    "mask_graph",
+    "spatial_neighbors",
+    "ligrec",
+    "centrality_scores",
+    "interaction_matrix",
+    "nhood_enrichment",
+    "calculate_niche",
+    "co_occurrence",
+    "spatial_autocorr",
+    "ripley",
+    "sepal",
+]

src/squidpy/gr/_build.py

@@ -6,7 +6,7 @@
 from collections.abc import Iterable  # noqa: F401
 from functools import partial
 from itertools import chain
-from typing import Any, cast
+from typing import Any, NamedTuple, cast
 
 import geopandas as gpd
 import numpy as np
@@ -61,6 +61,21 @@
 __all__ = ["spatial_neighbors"]
 
 
+class SpatialNeighborsResult(NamedTuple):
+    """Result of spatial_neighbors function.
+
+    Attributes
+    ----------
+    connectivities
+        Spatial connectivities matrix.
+    distances
+        Spatial distances matrix.
+    """
+
+    connectivities: csr_matrix
+    distances: csr_matrix
+
+
 @d.dedent
 @inject_docs(t=Transform, c=CoordType)
 def spatial_neighbors(
@@ -79,7 +94,7 @@ def spatial_neighbors(
     set_diag: bool = False,
     key_added: str = "spatial",
     copy: bool = False,
-) -> tuple[csr_matrix, csr_matrix] | None:
+) -> SpatialNeighborsResult | None:
     """
     Create a graph from spatial coordinates.
 
@@ -136,7 +151,7 @@ def spatial_neighbors(
 
     Returns
     -------
-    If ``copy = True``, returns a :class:`tuple` with the spatial connectivities and distances matrices.
+    If ``copy = True``, returns a :class:`~squidpy.gr.SpatialNeighborsResult` with the spatial connectivities and distances matrices.
 
     Otherwise, modifies the ``adata`` with the following keys:
 
@@ -259,7 +274,7 @@ def spatial_neighbors(
     }
 
     if copy:
-        return Adj, Dst
+        return SpatialNeighborsResult(connectivities=Adj, distances=Dst)
 
     _save_data(adata, attr="obsp", key=conns_key, data=Adj)
     _save_data(adata, attr="obsp", key=dists_key, data=Dst, prefix=False)

src/squidpy/gr/_nhood.py

@@ -4,14 +4,15 @@
 
 from collections.abc import Callable, Iterable, Sequence
 from functools import partial
-from typing import Any
+from typing import Any, NamedTuple
 
 import networkx as nx
 import numba.types as nt
 import numpy as np
 import pandas as pd
 from anndata import AnnData
 from numba import njit, prange  # noqa: F401
+from numpy.typing import NDArray
 from pandas import CategoricalDtype
 from scanpy import logging as logg
 from spatialdata import SpatialData
@@ -30,7 +31,24 @@
 
 __all__ = ["nhood_enrichment", "centrality_scores", "interaction_matrix"]
 
-dt = nt.uint32  # data type aliases (both for numpy and numba should match)
+
+class NhoodEnrichmentResult(NamedTuple):
+    """Result of nhood_enrichment function.
+
+    Attributes
+    ----------
+    zscore : NDArray[np.number]
+        Z-score values of enrichment statistic.
+    count : NDArray[np.number]
+        Enrichment count.
+    """
+
+    zscore: NDArray[np.number]
+    counts: NDArray[np.number]  # NamedTuple inherits from tuple so cannot use 'count' as attribute name
+
+
+# data type aliases (both for numpy and numba should match)
+dt = nt.uint32
 ndt = np.uint32
 _template = """
 @njit(dt[:, :](dt[:], dt[:], dt[:]), parallel={parallel}, fastmath=True)
@@ -131,7 +149,7 @@ def nhood_enrichment(
     n_jobs: int | None = None,
     backend: str = "loky",
     show_progress_bar: bool = True,
-) -> tuple[NDArrayA, NDArrayA] | None:
+) -> NhoodEnrichmentResult | None:
     """
     Compute neighborhood enrichment by permutation test.
 
@@ -149,7 +167,7 @@ def nhood_enrichment(
 
     Returns
     -------
-    If ``copy = True``, returns a :class:`tuple` with the z-score and the enrichment count.
+    If ``copy = True``, returns a :class:`~squidpy.gr.NhoodEnrichmentResult` with the z-score and the enrichment count.
 
     Otherwise, modifies the ``adata`` with the following keys:
 
@@ -202,7 +220,7 @@ def nhood_enrichment(
     zscore = (count - perms.mean(axis=0)) / perms.std(axis=0)
 
     if copy:
-        return zscore, count
+        return NhoodEnrichmentResult(zscore=zscore, counts=count)
 
     _save_data(
         adata,

src/squidpy/im/_io.py

@@ -25,7 +25,7 @@ def _assert_dims_present(dims: tuple[str, ...], include_z: bool = True) -> None:
 
 # modification of `skimage`'s `pil_to_ndarray`:
 # https://github.com/scikit-image/scikit-image/blob/main/skimage/io/_plugins/pil_plugin.py#L55
-def _infer_shape_dtype(fname: str) -> tuple[tuple[int, ...], np.dtype]:
+def _infer_shape_dtype(fname: str) -> tuple[tuple[int, ...], np.dtype[np.generic]]:
     def _palette_is_grayscale(pil_image: Image.Image) -> bool:
         # get palette as an array with R, G, B columns
         palette = np.asarray(pil_image.getpalette()).reshape((256, 3))
@@ -81,7 +81,7 @@ def _palette_is_grayscale(pil_image: Image.Image) -> bool:
     raise ValueError(f"Unable to infer image dtype for image mode `{image.mode}`.")
 
 
-def _get_image_shape_dtype(fname: str) -> tuple[tuple[int, ...], np.dtype]:
+def _get_image_shape_dtype(fname: str) -> tuple[tuple[int, ...], np.dtype[np.generic]]:
     try:
         return _infer_shape_dtype(fname)
     except Image.UnidentifiedImageError as e:
@@ -101,7 +101,7 @@ def _get_image_shape_dtype(fname: str) -> tuple[tuple[int, ...], np.dtype]:
 def _infer_dimensions(
     obj: NDArrayA | xr.DataArray | str,
     infer_dimensions: InferDimensions | tuple[str, ...] = InferDimensions.DEFAULT,
-) -> tuple[tuple[int, ...], tuple[str, ...], np.dtype, tuple[int, ...]]:
+) -> tuple[tuple[int, ...], tuple[str, ...], np.dtype[np.generic], tuple[int, ...]]:
     """
     Infer dimension names of an array.
 

tests/graph/test_nhood.py

@@ -49,13 +49,14 @@ def test_reproducibility(self, adata: AnnData, n_jobs: int):
         assert len(res1) == len(res2)
         assert len(res2) == len(res3)
 
-        for key in range(len(res1)):
-            np.testing.assert_array_equal(res2[key], res1[key])
-            if key == 0:  # z-score
-                with pytest.raises(AssertionError):
-                    np.testing.assert_array_equal(res3[key], res2[key])
-            else:  # counts
-                np.testing.assert_array_equal(res3[key], res2[key])
+        # Test that the same seed produces the same results
+        np.testing.assert_array_equal(res2.zscore, res1.zscore)
+        np.testing.assert_array_equal(res2.counts, res1.counts)
+
+        # Test that different seeds produce different z-scores but same counts
+        with pytest.raises(AssertionError):
+            np.testing.assert_array_equal(res3.zscore, res2.zscore)
+        np.testing.assert_array_equal(res3.counts, res2.counts)
 
 
 def test_centrality_scores(nhood_data: AnnData):

tests/graph/test_spatial_neighbors.py

@@ -198,25 +198,23 @@ def test_radius_min_max(self, non_visium_adata: AnnData, radius: tuple[float, fl
         np.testing.assert_allclose(spatial_dist, gt_ddist)
 
     def test_copy(self, non_visium_adata: AnnData):
-        conn, dist = spatial_neighbors(non_visium_adata, delaunay=True, coord_type=None, copy=True)
+        result = spatial_neighbors(non_visium_adata, delaunay=True, coord_type=None, copy=True)
 
-        assert isspmatrix_csr(conn)
-        assert isspmatrix_csr(dist)
+        assert isspmatrix_csr(result.connectivities)
+        assert isspmatrix_csr(result.distances)
         assert Key.obsp.spatial_conn() not in non_visium_adata.obsp
         assert Key.obsp.spatial_dist() not in non_visium_adata.obsp
-        np.testing.assert_allclose(dist.toarray(), self._gt_ddist)
-        np.testing.assert_allclose(conn.toarray(), self._gt_dgraph)
+        np.testing.assert_allclose(result.distances.toarray(), self._gt_ddist)
+        np.testing.assert_allclose(result.connectivities.toarray(), self._gt_dgraph)
 
     @pytest.mark.parametrize("percentile", [99.0, 95.0])
     def test_percentile_filtering(self, adata_hne: AnnData, percentile: float, coord_type="generic"):
-        conn, dist = spatial_neighbors(adata_hne, coord_type=coord_type, copy=True)
-        conn_filtered, dist_filtered = spatial_neighbors(
-            adata_hne, coord_type=coord_type, percentile=percentile, copy=True
-        )
+        result = spatial_neighbors(adata_hne, coord_type=coord_type, copy=True)
+        result_filtered = spatial_neighbors(adata_hne, coord_type=coord_type, percentile=percentile, copy=True)
 
         # check whether there are less connectivities in the filtered graph and whether the max distance is smaller
-        assert not ((conn != conn_filtered).nnz == 0)
-        assert dist.max() > dist_filtered.max()
+        assert not ((result.connectivities != result_filtered.connectivities).nnz == 0)
+        assert result.distances.max() > result_filtered.distances.max()
 
         Adj, Dst = _build_connectivity(adata_hne.obsm["spatial"], n_neighs=6, return_distance=True, set_diag=False)
         threshold = np.percentile(Dst.data, percentile)
@@ -225,7 +223,7 @@ def test_percentile_filtering(self, adata_hne: AnnData, percentile: float, coord
         Adj.eliminate_zeros()
         Dst.eliminate_zeros()
 
-        assert dist_filtered.max() == Dst.max()
+        assert result_filtered.distances.max() == Dst.max()
 
     @pytest.mark.parametrize("n_neighs", [5, 10, 20])
     def test_spatial_neighbors_generic(self, n_neighs: int):