docs: clean up typing module (#2325)

Author: flying-sheep

docs/api.md

@@ -253,12 +253,17 @@ Types used by the former:
    abc.CSCDataset
 ```
 
-.. these are types, not classes
+<!-- these are types, not classes, so don’t use the above template -->
 
 ```{eval-rst}
+.. toctree::
+   :hidden:
+
+   typing
+
 .. autosummary::
-   :toctree: generated/
 
+   typing.Index1D
    typing.Index
    typing.AxisStorable
    typing.RWAble

docs/conf.py

@@ -145,9 +145,10 @@ def res(
         for kind in ["", "View"]
     },
     # Can’t use `set_module` for `type`s. When moving out of .experimental, define in actual location.
-    "anndata.compat.Index": "anndata.typing.Index",
     "anndata._types.StorageType": "anndata.experimental.StorageType",
     # https://github.com/theislab/scanpydoc/issues/254
+    "anndata.typing.Index1D": "anndata.typing.Index1D",
+    "anndata.typing.Index": "anndata.typing.Index",
     "anndata.typing.RWAble": "anndata.typing.RWAble",
     "anndata.typing.AxisStorable": "anndata.typing.AxisStorable",
     #### h5py

docs/typing.md

@@ -0,0 +1,10 @@
+# anndata.typing
+
+```{eval-rst}
+.. module:: anndata.typing
+
+.. autotype:: Index1D
+.. autotype:: Index
+.. autotype:: AxisStorable
+.. autotype:: RWAble
+```

src/anndata/_core/anndata.py

@@ -61,8 +61,8 @@
 
     from zarr.storage import StoreLike
 
-    from ..compat import Index1D, Index1DNorm, XDataset
-    from ..typing import XDataType
+    from ..compat import XDataset
+    from ..typing import Index1D, _Index1DNorm, _XDataType
     from .aligned_mapping import AxisArraysView, LayersView, PairwiseArraysView
     from .index import Index
 
@@ -206,8 +206,8 @@ class AnnData(metaclass=utils.DeprecationMixinMeta):  # noqa: PLW1641
 
     # view attributes
     _adata_ref: AnnData | None
-    _oidx: Index1DNorm | None
-    _vidx: Index1DNorm | None
+    _oidx: _Index1DNorm | None
+    _vidx: _Index1DNorm | None
 
     @old_positionals(
         "obsm",
@@ -222,14 +222,14 @@ class AnnData(metaclass=utils.DeprecationMixinMeta):  # noqa: PLW1641
     )
     def __init__(  # noqa: PLR0913
         self,
-        X: XDataType | pd.DataFrame | None = None,
+        X: _XDataType | pd.DataFrame | None = None,
         obs: pd.DataFrame | Mapping[str, Iterable[Any]] | None = None,
         var: pd.DataFrame | Mapping[str, Iterable[Any]] | None = None,
         uns: Mapping[str, Any] | None = None,
         *,
         obsm: np.ndarray | Mapping[str, Sequence[Any]] | None = None,
         varm: np.ndarray | Mapping[str, Sequence[Any]] | None = None,
-        layers: Mapping[str, XDataType] | None = None,
+        layers: Mapping[str, _XDataType] | None = None,
         raw: Mapping[str, Any] | None = None,
         dtype: np.dtype | type | str | None = None,
         shape: tuple[int, int] | None = None,
@@ -238,8 +238,8 @@ def __init__(  # noqa: PLR0913
         asview: bool = False,
         obsp: np.ndarray | Mapping[str, Sequence[Any]] | None = None,
         varp: np.ndarray | Mapping[str, Sequence[Any]] | None = None,
-        oidx: Index1DNorm | int | np.integer | None = None,
-        vidx: Index1DNorm | int | np.integer | None = None,
+        oidx: _Index1DNorm | int | np.integer | None = None,
+        vidx: _Index1DNorm | int | np.integer | None = None,
     ):
         # check for any multi-indices that aren’t later checked in coerce_array
         for attr, key in [(obs, "obs"), (var, "var"), (X, "X")]:
@@ -273,8 +273,8 @@ def __init__(  # noqa: PLR0913
     def _init_as_view(
         self,
         adata_ref: AnnData,
-        oidx: Index1DNorm | int | np.integer,
-        vidx: Index1DNorm | int | np.integer,
+        oidx: _Index1DNorm | int | np.integer,
+        vidx: _Index1DNorm | int | np.integer,
     ):
         if adata_ref.isbacked and adata_ref.is_view:
             msg = (
@@ -571,7 +571,7 @@ def shape(self) -> tuple[int, int]:
         return self.n_obs, self.n_vars
 
     @property
-    def X(self) -> XDataType | None:
+    def X(self) -> _XDataType | None:
         """Data matrix of shape :attr:`n_obs` × :attr:`n_vars`."""
         if self.isbacked:
             if not self.file.is_open:
@@ -602,7 +602,7 @@ def X(self) -> XDataType | None:
         #     return X
 
     @X.setter
-    def X(self, value: XDataType | None):  # noqa: PLR0912
+    def X(self, value: _XDataType | None):  # noqa: PLR0912
         if value is None:
             if self.isbacked:
                 msg = "Cannot currently remove data matrix from backed object."
@@ -1066,7 +1066,7 @@ def _set_backed(self, attr, value):
 
     def _normalize_indices(
         self, index: Index | None
-    ) -> tuple[Index1DNorm | int | np.integer, Index1DNorm | int | np.integer]:
+    ) -> tuple[_Index1DNorm | int | np.integer, _Index1DNorm | int | np.integer]:
         return _normalize_indices(index, self.obs_names, self.var_names)
 
     # TODO: this is not quite complete...
@@ -1244,7 +1244,7 @@ def _inplace_subset_obs(self, index: Index1D):
         self._init_as_actual(adata_subset)
 
     # TODO: Update, possibly remove
-    def __setitem__(self, index: Index, val: float | XDataType):
+    def __setitem__(self, index: Index, val: float | _XDataType):
         if self.is_view:
             msg = "Object is view and cannot be accessed with `[]`."
             raise ValueError(msg)

src/anndata/_core/file_backing.py

@@ -17,7 +17,7 @@
     from os import PathLike
     from typing import Literal
 
-    from .._types import ArrayStorageType
+    from .._types import _ArrayStorageType
     from . import anndata
 
 
@@ -143,7 +143,7 @@ def to_memory(x, *, copy: bool = False):
 
 @to_memory.register(ZarrArray)
 @to_memory.register(h5py.Dataset)
-def _(x: ArrayStorageType, *, copy: bool = False):
+def _(x: _ArrayStorageType, *, copy: bool = False):
     return x[...]
 
 

src/anndata/_core/index.py

@@ -16,12 +16,12 @@
 if TYPE_CHECKING:
     from numpy.typing import NDArray
 
-    from ..compat import Index, Index1D, Index1DNorm
+    from ..typing import Index, Index1D, _Index1DNorm
 
 
 def _normalize_indices(
     index: Index | None, names0: pd.Index, names1: pd.Index
-) -> tuple[Index1DNorm | int | np.integer, Index1DNorm | int | np.integer]:
+) -> tuple[_Index1DNorm | int | np.integer, _Index1DNorm | int | np.integer]:
     # deal with tuples of length 1
     if isinstance(index, tuple) and len(index) == 1:
         index = index[0]
@@ -33,7 +33,7 @@ def _normalize_indices(
 
 def _normalize_index(  # noqa: PLR0911, PLR0912
     indexer: Index1D, index: pd.Index
-) -> Index1DNorm | int | np.integer:
+) -> _Index1DNorm | int | np.integer:
     # TODO: why is this here? All tests pass without it and it seems at the minimum not strict enough.
     if not isinstance(index, pd.RangeIndex) and index.dtype in (np.float64, np.int64):
         msg = f"Don’t call _normalize_index with non-categorical/string names and non-range index {index}"
@@ -172,7 +172,7 @@ def unpack_index(index: Index) -> tuple[Index1D, Index1D]:
 @singledispatch
 def _subset(
     a: np.ndarray | pd.DataFrame,
-    subset_idx: tuple[Index1DNorm] | tuple[Index1DNorm, Index1DNorm],
+    subset_idx: tuple[_Index1DNorm] | tuple[_Index1DNorm, _Index1DNorm],
 ):
     # Select as combination of indexes, not coordinates
     # Correcting for indexing behaviour of np.ndarray
@@ -183,7 +183,7 @@ def _subset(
 
 @_subset.register(DaskArray)
 def _subset_dask(
-    a: DaskArray, subset_idx: tuple[Index1DNorm] | tuple[Index1DNorm, Index1DNorm]
+    a: DaskArray, subset_idx: tuple[_Index1DNorm] | tuple[_Index1DNorm, _Index1DNorm]
 ):
     if len(subset_idx) > 1 and all(isinstance(x, Iterable) for x in subset_idx):
         if issparse(a._meta) and a._meta.format == "csc":
@@ -196,7 +196,7 @@ def _subset_dask(
 @_subset.register(CSArray)
 def _subset_sparse(
     a: CSMatrix | CSArray,
-    subset_idx: tuple[Index1DNorm] | tuple[Index1DNorm, Index1DNorm],
+    subset_idx: tuple[_Index1DNorm] | tuple[_Index1DNorm, _Index1DNorm],
 ):
     # Correcting for indexing behaviour of sparse.spmatrix
     if len(subset_idx) > 1 and all(isinstance(x, Iterable) for x in subset_idx):
@@ -211,14 +211,14 @@ def _subset_sparse(
 @_subset.register(Dataset2D)
 def _subset_df(
     df: pd.DataFrame | Dataset2D,
-    subset_idx: tuple[Index1DNorm] | tuple[Index1DNorm, Index1DNorm],
+    subset_idx: tuple[_Index1DNorm] | tuple[_Index1DNorm, _Index1DNorm],
 ):
     return df.iloc[subset_idx]
 
 
 @_subset.register(AwkArray)
 def _subset_awkarray(
-    a: AwkArray, subset_idx: tuple[Index1DNorm] | tuple[Index1DNorm, Index1DNorm]
+    a: AwkArray, subset_idx: tuple[_Index1DNorm] | tuple[_Index1DNorm, _Index1DNorm]
 ):
     if all(isinstance(x, Iterable) for x in subset_idx):
         subset_idx = np.ix_(*subset_idx)
@@ -228,7 +228,7 @@ def _subset_awkarray(
 # Registration for SparseDataset occurs in sparse_dataset.py
 @_subset.register(h5py.Dataset)
 def _subset_dataset(
-    d: h5py.Dataset, subset_idx: tuple[Index1DNorm] | tuple[Index1DNorm, Index1DNorm]
+    d: h5py.Dataset, subset_idx: tuple[_Index1DNorm] | tuple[_Index1DNorm, _Index1DNorm]
 ):
     order: tuple[NDArray[np.integer] | slice, ...]
     inv_order: tuple[NDArray[np.integer] | slice, ...]
@@ -252,8 +252,8 @@ def _index_order_and_inverse(
 @overload
 def _index_order_and_inverse(axis_idx: slice) -> tuple[slice, slice]: ...
 def _index_order_and_inverse(
-    axis_idx: Index1DNorm,
-) -> tuple[Index1DNorm, NDArray[np.integer] | slice]:
+    axis_idx: _Index1DNorm,
+) -> tuple[_Index1DNorm, NDArray[np.integer] | slice]:
     """Order and get inverse index array."""
     if not isinstance(axis_idx, np.ndarray):
         return axis_idx, slice(None)
@@ -270,8 +270,8 @@ def _process_index_for_h5py(
 @overload
 def _process_index_for_h5py(idx: slice) -> tuple[slice, None]: ...
 def _process_index_for_h5py(
-    idx: Index1DNorm,
-) -> tuple[Index1DNorm, NDArray[np.integer] | None]:
+    idx: _Index1DNorm,
+) -> tuple[_Index1DNorm, NDArray[np.integer] | None]:
     """Process a single index for h5py compatibility, handling sorting and duplicates."""
     if not isinstance(idx, np.ndarray):
         # Not an array (slice, integer, list) - no special processing needed
@@ -294,7 +294,7 @@ def _process_index_for_h5py(
 
 def _safe_fancy_index_h5py(
     dataset: h5py.Dataset,
-    subset_idx: tuple[Index1DNorm] | tuple[Index1DNorm, Index1DNorm],
+    subset_idx: tuple[_Index1DNorm] | tuple[_Index1DNorm, _Index1DNorm],
 ) -> h5py.Dataset:
     # Handle multi-dimensional indexing of h5py dataset
     # This avoids h5py's limitation with multi-dimensional fancy indexing
@@ -331,7 +331,7 @@ def _safe_fancy_index_h5py(
     return result
 
 
-def _get_index_size(idx: Index1DNorm, dim_size: int) -> int:
+def _get_index_size(idx: _Index1DNorm, dim_size: int) -> int:
     """Get size for any index type."""
     if isinstance(idx, slice):
         return len(range(*idx.indices(dim_size)))

src/anndata/_core/raw.py

@@ -17,7 +17,8 @@
     from collections.abc import Mapping, Sequence
     from typing import ClassVar
 
-    from ..compat import CSMatrix, Index, Index1DNorm
+    from ..compat import CSMatrix
+    from ..typing import Index, _Index1DNorm
     from .aligned_mapping import AxisArraysView
     from .anndata import AnnData
     from .sparse_dataset import BaseCompressedSparseDataset
@@ -171,7 +172,7 @@ def to_adata(self) -> AnnData:
 
     def _normalize_indices(
         self, packed_index: Index
-    ) -> tuple[Index1DNorm | int | np.integer, Index1DNorm | int | np.integer]:
+    ) -> tuple[_Index1DNorm | int | np.integer, _Index1DNorm | int | np.integer]:
         # deal with slicing with pd.Series
         if isinstance(packed_index, pd.Series):
             packed_index = packed_index.values

src/anndata/_core/sparse_dataset.py

@@ -48,9 +48,8 @@
     from types import ModuleType
     from typing import Any, Literal
 
-    from .._types import ArrayStorageType, GroupStorageType
-    from ..compat import Index1DNorm
-    from .index import Index, Index1D
+    from .._types import _ArrayStorageType, _GroupStorageType
+    from ..typing import Index, Index1D, _Index1DNorm
 
 
 type DenseType = np.ndarray | CupyArray
@@ -97,7 +96,7 @@ def slice_as_int(s: slice, l: int) -> int:
 
 
 @dataclass
-class BackedSparseMatrix[ArrayT: ArrayStorageType]:
+class BackedSparseMatrix[ArrayT: _ArrayStorageType]:
     """\
     Mixin class for backed sparse matrices.
 
@@ -310,7 +309,7 @@ def mean_slice_length(slices) -> int:
         )
 
 
-def _get_group_format(group: GroupStorageType) -> str:
+def _get_group_format(group: _GroupStorageType) -> str:
     if "h5sparse_format" in group.attrs:
         # TODO: Warn about an old format
         # If this is only just going to be public, I could insist it's not like this
@@ -331,7 +330,7 @@ def is_sparse_indexing_overridden(
     )
 
 
-class BaseCompressedSparseDataset[GroupT: GroupStorageType, ArrayT: ArrayStorageType](
+class BaseCompressedSparseDataset[GroupT: _GroupStorageType, ArrayT: _ArrayStorageType](
     abc._AbstractCSDataset, ABC
 ):
     _group: GroupT
@@ -597,7 +596,7 @@ class _CSCDataset(BaseCompressedSparseDataset, abc.CSCDataset):
 
 @doctest_filterwarnings("ignore", r"Moving element.*uns.*to.*obsp", FutureWarning)
 def sparse_dataset(
-    group: GroupStorageType,
+    group: _GroupStorageType,
     *,
     should_cache_indptr: bool = True,
 ) -> abc.CSRDataset | abc.CSCDataset:
@@ -668,6 +667,6 @@ def sparse_dataset(
 
 @_subset.register(BaseCompressedSparseDataset)
 def subset_sparsedataset(
-    d, subset_idx: tuple[Index1DNorm] | tuple[Index1DNorm, Index1DNorm]
+    d, subset_idx: tuple[_Index1DNorm] | tuple[_Index1DNorm, _Index1DNorm]
 ):
     return d[subset_idx]

src/anndata/_core/storage.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, get_args
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -12,6 +12,7 @@
 from ..compat import XDataset
 from ..utils import (
     ensure_df_homogeneous,
+    get_union_members,
     join_english,
     raise_value_error_if_multiindex_columns,
     warn,
@@ -30,13 +31,13 @@ def coerce_array(
     allow_array_like: bool = False,
 ):
     """Coerce arrays stored in layers/X, and aligned arrays ({obs,var}{m,p})."""
-    from ..typing import ArrayDataStructureTypes
+    from ..typing import _ArrayDataStructureTypes
 
     # If value is a scalar and we allow that, return it
     if allow_array_like and np.isscalar(value):
         return value
     # If value is one of the allowed types, return it
-    array_data_structure_types = get_args(ArrayDataStructureTypes)
+    array_data_structure_types = get_union_members(_ArrayDataStructureTypes)
     if isinstance(value, XDataset):
         value = Dataset2D(value)
     if isinstance(value, (*array_data_structure_types, Dataset2D)):