Add overloads to docs (#56)

Author: flying-sheep

.pre-commit-config.yaml

@@ -34,5 +34,7 @@ repos:
           - zarr
           - h5py
           - anndata
+          - types-docutils
+          - sphinx
 ci:
   skip: [mypy]  # too big

docs/conf.py

@@ -6,6 +6,14 @@
 from datetime import UTC, datetime
 from importlib.metadata import metadata
 from pathlib import Path
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from docutils.nodes import TextElement, reference
+    from sphinx.addnodes import pending_xref
+    from sphinx.application import Sphinx
+    from sphinx.environment import BuildEnvironment
 
 
 HERE = Path(__file__).parent
@@ -47,6 +55,7 @@
 napoleon_numpy_docstring = True
 todo_include_todos = False
 intersphinx_mapping = dict(
+    anndata=("https://anndata.readthedocs.io/en/stable/", None),
     cupy=("https://docs.cupy.dev/en/stable/", None),
     dask=("https://docs.dask.org/en/stable/", None),
     h5py=("https://docs.h5py.org/en/stable/", None),
@@ -55,47 +64,14 @@
     scipy=("https://docs.scipy.org/doc/scipy/", None),
     zarr=("https://zarr.readthedocs.io/en/stable/", None),
 )
-# Try overriding type paths
-qualname_overrides = autodoc_type_aliases = {
-    "np.bool": ("py:data", "numpy.bool"),
-    "np.dtype": "numpy.dtype",
-    "np.number": "numpy.number",
-    "np.integer": "numpy.integer",
-    "np.floating": "numpy.floating",
-    "np.random.Generator": "numpy.random.Generator",
-    "ArrayLike": "numpy.typing.ArrayLike",
-    "DTypeLike": "numpy.typing.DTypeLike",
-    "NDArray": "numpy.typing.NDArray",
-    "_pytest.fixtures.FixtureRequest": "pytest.FixtureRequest",
-    **{
-        k: v
-        for k_plain, v in {
-            "CSBase": "scipy.sparse.spmatrix",
-            "CupyArray": "cupy.ndarray",
-            "CupySparseMatrix": "cupyx.scipy.sparse.spmatrix",
-            "DaskArray": "dask.array.Array",
-            "H5Dataset": "h5py.Dataset",
-            "ZarrArray": "zarr.Array",
-        }.items()
-        for k in (k_plain, f"types.{k_plain}")
-    },
-}
-# If that doesn’t work, ignore them
-nitpick_ignore = {
-    ("py:class", "fast_array_utils.types.T_co"),
+nitpick_ignore = [
     ("py:class", "Arr"),
+    ("py:class", "ToDType"),
     ("py:class", "testing.fast_array_utils._array_type.Arr"),
     ("py:class", "testing.fast_array_utils._array_type.Inner"),
     ("py:class", "_DTypeLikeFloat32"),
     ("py:class", "_DTypeLikeFloat64"),
-    # sphinx bugs, should be covered by `autodoc_type_aliases` above
-    ("py:class", "Array"),
-    ("py:class", "ArrayLike"),
-    ("py:class", "DTypeLike"),
-    ("py:class", "NDArray"),
-    ("py:class", "np.bool"),
-    ("py:class", "_pytest.fixtures.FixtureRequest"),
-}
+]
 
 # Options for HTML output
 html_theme = "furo"
@@ -104,3 +80,62 @@
     source_branch="main",
     source_directory="docs/",
 )
+
+_np_nocls = {"float64": "attr"}
+_optional_types = {
+    "CupyArray": "cupy.ndarray",
+    "CupySparseMatrix": "cupyx.scipy.sparse.spmatrix",
+    "DaskArray": "dask.array.Array",
+    "H5Dataset": "h5py.Dataset",
+    "ZarrArray": "zarr.Array",
+}
+
+
+def find_type_alias(name: str) -> tuple[str, str] | tuple[None, None]:
+    """Find a type alias."""
+    import numpy.typing as npt
+
+    from fast_array_utils import types, typing
+
+    if name in typing.__all__:
+        return "data", f"fast_array_utils.typing.{name}"
+    if name.startswith("types.") and name[6:] in types.__all__:
+        if path := _optional_types.get(name[6:]):
+            return "class", path
+        return "data", f"fast_array_utils.{name}"
+    if name.startswith("np."):
+        return _np_nocls.get(name[3:], "class"), f"numpy.{name[3:]}"
+    if name in npt.__all__:
+        return "data", f"numpy.typing.{name}"
+    return None, None
+
+
+def resolve_type_aliases(
+    app: Sphinx, env: BuildEnvironment, node: pending_xref, contnode: TextElement
+) -> reference | None:
+    """Resolve :class: references to our type aliases as :attr: instead."""
+    if (node["refdomain"], node["reftype"]) != ("py", "class"):
+        return None
+    typ, target = find_type_alias(node["reftarget"])
+    if typ is None or target is None:
+        return None
+    if target.startswith("fast_array_utils."):
+        ref = env.get_domain("py").resolve_xref(
+            env, node["refdoc"], app.builder, typ, target, node, contnode
+        )
+    else:
+        from sphinx.ext.intersphinx import resolve_reference_any_inventory
+
+        node["reftype"] = typ
+        node["reftarget"] = target
+        ref = resolve_reference_any_inventory(
+            env=env, honor_disabled_refs=False, node=node, contnode=contnode
+        )
+    if ref is None:
+        msg = f"Could not resolve {typ} {target} (from {node['reftarget']})"
+        raise AssertionError(msg)
+    return ref
+
+
+def setup(app: Sphinx) -> None:  # noqa: D103
+    app.connect("missing-reference", resolve_type_aliases, priority=800)

docs/index.rst

@@ -10,6 +10,18 @@
 .. automodule:: fast_array_utils
    :members:
 
+``fast_array_utils.types``
+--------------------------
+
+.. automodule:: fast_array_utils.types
+   :members: CSBase, CSDataset
+
+``fast_array_utils.typing``
+---------------------------
+
+.. automodule:: fast_array_utils.typing
+   :members:
+
 ``fast_array_utils.conv``
 -------------------------
 

src/fast_array_utils/conv/__init__.py

@@ -3,7 +3,59 @@
 
 from __future__ import annotations
 
-from ._to_dense import to_dense
+from typing import TYPE_CHECKING, overload
+
+from ..typing import CpuArray, DiskArray, GpuArray  # noqa: TC001
+from ._to_dense import to_dense_
+
+
+if TYPE_CHECKING:
+    from typing import Any, Literal
+
+    from numpy.typing import NDArray
+
+    from .. import types
 
 
 __all__ = ["to_dense"]
+
+
+@overload
+def to_dense(
+    x: CpuArray | DiskArray | types.CSDataset, /, *, to_memory: bool = False
+) -> NDArray[Any]: ...
+
+
+@overload
+def to_dense(x: types.DaskArray, /, *, to_memory: Literal[False] = False) -> types.DaskArray: ...
+@overload
+def to_dense(x: types.DaskArray, /, *, to_memory: Literal[True]) -> NDArray[Any]: ...
+
+
+@overload
+def to_dense(x: GpuArray, /, *, to_memory: Literal[False] = False) -> types.CupyArray: ...
+@overload
+def to_dense(x: GpuArray, /, *, to_memory: Literal[True]) -> NDArray[Any]: ...
+
+
+def to_dense(
+    x: CpuArray | GpuArray | DiskArray | types.CSDataset | types.DaskArray,
+    /,
+    *,
+    to_memory: bool = False,
+) -> NDArray[Any] | types.DaskArray | types.CupyArray:
+    """Convert x to a dense array.
+
+    Parameters
+    ----------
+    x
+        Input object to be converted.
+    to_memory
+        Also load data into memory (resulting in a :class:`numpy.ndarray`).
+
+    Returns
+    -------
+    Dense form of ``x``
+
+    """
+    return to_dense_(x, to_memory=to_memory)

src/fast_array_utils/conv/_to_dense.py

@@ -2,106 +2,61 @@
 from __future__ import annotations
 
 from functools import singledispatch
-from typing import TYPE_CHECKING, cast, overload
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 
 from .. import types
+from ..typing import CpuArray, DiskArray, GpuArray  # noqa: TC001
 
 
 if TYPE_CHECKING:
-    from typing import Any, Literal, TypeAlias
+    from typing import Any
 
     from numpy.typing import NDArray
 
-    MemDiskArray: TypeAlias = (
-        NDArray[Any] | types.CSBase | types.H5Dataset | types.ZarrArray | types.CSDataset
-    )
-    Array: TypeAlias = MemDiskArray | types.CupyArray | types.CupySparseMatrix | types.DaskArray
-
-
-__all__ = ["to_dense"]
-
-
-@overload
-def to_dense(x: MemDiskArray, /, *, to_memory: bool = False) -> NDArray[Any]: ...
-
-
-@overload
-def to_dense(x: types.DaskArray, /, *, to_memory: Literal[False] = False) -> types.DaskArray: ...
-@overload
-def to_dense(x: types.DaskArray, /, *, to_memory: Literal[True]) -> NDArray[Any]: ...
-
-
-@overload
-def to_dense(
-    x: types.CupyArray | types.CupySparseMatrix, /, *, to_memory: Literal[False] = False
-) -> types.CupyArray: ...
-@overload
-def to_dense(
-    x: types.CupyArray | types.CupySparseMatrix, /, *, to_memory: Literal[True]
-) -> NDArray[Any]: ...
-
-
-def to_dense(
-    x: Array, /, *, to_memory: bool = False
-) -> NDArray[Any] | types.DaskArray | types.CupyArray:
-    """Convert x to a dense array.
-
-    Parameters
-    ----------
-    x
-        Input object to be converted.
-    to_memory
-        Also load data into memory (resulting in a :class:`numpy.ndarray`).
-
-    Returns
-    -------
-    Dense form of ``x``
-
-    """
-    return _to_dense(x, to_memory=to_memory)
-
 
 # fallback’s arg0 type has to include types of registered functions
 @singledispatch
-def _to_dense(
-    x: Array, /, *, to_memory: bool = False
-) -> NDArray[Any] | types.DaskArray | types.CupyArray:
+def to_dense_(
+    x: CpuArray | GpuArray | DiskArray | types.DaskArray, /, *, to_memory: bool = False
+) -> NDArray[Any] | types.CupyArray | types.DaskArray:
     del to_memory  # it already is
     return np.asarray(x)
 
 
-@_to_dense.register(types.CSBase)  # type: ignore[call-overload,misc]
+@to_dense_.register(types.CSBase)  # type: ignore[call-overload,misc]
 def _to_dense_cs(x: types.CSBase, /, *, to_memory: bool = False) -> NDArray[Any]:
     from . import scipy
 
     del to_memory  # it already is
     return scipy.to_dense(x)
 
 
-@_to_dense.register(types.DaskArray)
+@to_dense_.register(types.DaskArray)
 def _to_dense_dask(
     x: types.DaskArray, /, *, to_memory: bool = False
 ) -> NDArray[Any] | types.DaskArray:
     import dask.array as da
 
+    from . import to_dense
+
     x = da.map_blocks(to_dense, x)
     return x.compute() if to_memory else x  # type: ignore[return-value]
 
 
-@_to_dense.register(types.CSDataset)
+@to_dense_.register(types.CSDataset)
 def _to_dense_ooc(x: types.CSDataset, /, *, to_memory: bool = False) -> NDArray[Any]:
+    from . import to_dense
+
     if not to_memory:
         msg = "to_memory must be True if x is an CS{R,C}Dataset"
         raise ValueError(msg)
     # TODO(flying-sheep): why is to_memory of type Any?  # noqa: TD003
     return to_dense(cast("types.CSBase", x.to_memory()))
 
 
-@_to_dense.register(types.CupyArray | types.CupySparseMatrix)  # type: ignore[call-overload,misc]
-def _to_dense_cupy(
-    x: types.CupyArray | types.CupySparseMatrix, /, *, to_memory: bool = False
-) -> NDArray[Any] | types.CupyArray:
+@to_dense_.register(GpuArray)  # type: ignore[call-overload,misc]
+def _to_dense_cupy(x: GpuArray, /, *, to_memory: bool = False) -> NDArray[Any] | types.CupyArray:
     x = x.toarray() if isinstance(x, types.CupySparseMatrix) else x
     return x.get() if to_memory else x