gh-355: consistent use of typing and collections.abc (#356)

Fixes #355, simplifying #308, xref #350. Change all import of
`typing`/`collections.abc` to `import <x>` rather than `from <x> import`
syntax. This is useful because it is straightforward to see at a glance
where the `import` is coming from. Further, we have both
`numpy.random.Generator` and `collections.abc.Generator` in use - so it
separates them. I've also moved everything out of the `TYPE_CHECKING`
block, except from where `ruff` says we should.

Author: paddyroddy

glass/core/algorithm.py

@@ -2,12 +2,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 import numpy as np
-
-if TYPE_CHECKING:
-    import numpy.typing as npt
+import numpy.typing as npt
 
 
 def nnls(

glass/fields.py

@@ -26,28 +26,30 @@
 
 from __future__ import annotations
 
+import collections.abc
+import typing
 import warnings
-from collections.abc import Generator, Iterable, Sequence
-from typing import Any, Callable, Optional, Union
 
 import healpy as hp
 import numpy as np
 import numpy.typing as npt
 from gaussiancl import gaussiancl
 
 # types
-Size = Optional[Union[int, tuple[int, ...]]]
-Iternorm = tuple[Optional[int], npt.NDArray, npt.NDArray]
-ClTransform = Union[str, Callable[[npt.NDArray], npt.NDArray]]
-Cls = Sequence[Union[npt.NDArray, Sequence[float]]]
+Size = typing.Optional[typing.Union[int, tuple[int, ...]]]
+Iternorm = tuple[typing.Optional[int], npt.NDArray, npt.NDArray]
+ClTransform = typing.Union[str, typing.Callable[[npt.NDArray], npt.NDArray]]
+Cls = collections.abc.Sequence[
+    typing.Union[npt.NDArray, collections.abc.Sequence[float]]
+]
 Alms = npt.NDArray
 
 
 def iternorm(
     k: int,
-    cov: Iterable[npt.NDArray],
+    cov: collections.abc.Iterable[npt.NDArray],
     size: Size = None,
-) -> Generator[Iternorm, None, None]:
+) -> collections.abc.Generator[Iternorm, None, None]:
     """Return the vector a and variance sigma^2 for iterative normal sampling."""
     n: tuple[int, ...]
     if size is None:
@@ -105,7 +107,9 @@ def iternorm(
         yield j, a, s
 
 
-def cls2cov(cls: Cls, nl: int, nf: int, nc: int) -> Generator[npt.NDArray, None, None]:
+def cls2cov(
+    cls: Cls, nl: int, nf: int, nc: int
+) -> collections.abc.Generator[npt.NDArray, None, None]:
     """Return array of cls as a covariance matrix for iterative sampling."""
     cov = np.zeros((nl, nc + 1))
     end = 0
@@ -134,7 +138,7 @@ def multalm(alm: Alms, bl: npt.NDArray, *, inplace: bool = False) -> Alms:
     return out
 
 
-def transform_cls(cls: Cls, tfm: ClTransform, pars: tuple[Any, ...] = ()) -> Cls:
+def transform_cls(cls: Cls, tfm: ClTransform, pars: tuple[typing.Any, ...] = ()) -> Cls:
     """Transform Cls to Gaussian Cls."""
     gls = []
     for cl in cls:
@@ -212,7 +216,7 @@ def generate_gaussian(
     *,
     ncorr: int | None = None,
     rng: np.random.Generator | None = None,
-) -> Generator[npt.NDArray, None, None]:
+) -> collections.abc.Generator[npt.NDArray, None, None]:
     """
     Sample Gaussian random fields from Cls iteratively.
 
@@ -298,7 +302,7 @@ def generate_lognormal(
     *,
     ncorr: int | None = None,
     rng: np.random.Generator | None = None,
-) -> Generator[npt.NDArray, None, None]:
+) -> collections.abc.Generator[npt.NDArray, None, None]:
     """Sample lognormal random fields from Gaussian Cls iteratively."""
     for i, m in enumerate(generate_gaussian(gls, nside, ncorr=ncorr, rng=rng)):
         # compute the variance of the auto-correlation

glass/galaxies.py

@@ -19,19 +19,18 @@
 
 from __future__ import annotations
 
+import typing
 import warnings
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import numpy.typing as npt
-
-    from glass.shells import RadialWindow
 
 import healpix
 import numpy as np
+import numpy.typing as npt
 
 from glass.core.array import broadcast_leading_axes, cumtrapz
 
+if typing.TYPE_CHECKING:
+    from glass.shells import RadialWindow
+
 
 def redshifts(
     n: int | npt.ArrayLike,

glass/lensing.py

@@ -31,15 +31,14 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import typing
 
 import healpy as hp
 import numpy as np
+import numpy.typing as npt
 
-if TYPE_CHECKING:
-    from collections.abc import Sequence
-
-    import numpy.typing as npt
+if typing.TYPE_CHECKING:
+    import collections.abc
 
     from cosmology import Cosmology
 
@@ -363,7 +362,7 @@ def wlens(self) -> float:
 
 
 def multi_plane_matrix(
-    shells: Sequence[RadialWindow],
+    shells: collections.abc.Sequence[RadialWindow],
     cosmo: Cosmology,
 ) -> npt.ArrayLike:
     """Compute the matrix of lensing contributions from each shell."""
@@ -377,7 +376,7 @@ def multi_plane_matrix(
 
 def multi_plane_weights(
     weights: npt.ArrayLike,
-    shells: Sequence[RadialWindow],
+    shells: collections.abc.Sequence[RadialWindow],
     cosmo: Cosmology,
 ) -> npt.ArrayLike:
     """

glass/observations.py

@@ -29,16 +29,13 @@
 from __future__ import annotations
 
 import math
-from typing import TYPE_CHECKING
 
 import healpy as hp
 import numpy as np
+import numpy.typing as npt
 
 from glass.core.array import cumtrapz
 
-if TYPE_CHECKING:
-    import numpy.typing as npt
-
 
 def vmap_galactic_ecliptic(
     nside: int,

glass/shapes.py

@@ -24,12 +24,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 import numpy as np
-
-if TYPE_CHECKING:
-    import numpy.typing as npt
+import numpy.typing as npt
 
 
 def triaxial_axis_ratio(zeta, xi, size=None, *, rng=None):

glass/shells.py

@@ -44,21 +44,21 @@
 
 from __future__ import annotations
 
+import collections.abc
+import typing
 import warnings
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Callable, NamedTuple, Union
 
 import numpy as np
 import numpy.typing as npt
 
 from glass.core.array import ndinterp
 
-if TYPE_CHECKING:
+if typing.TYPE_CHECKING:
     from cosmology import Cosmology
 
 # types
-ArrayLike1D = Union[Sequence[float], npt.NDArray]
-WeightFunc = Callable[[ArrayLike1D], npt.NDArray]
+ArrayLike1D = typing.Union[collections.abc.Sequence[float], npt.NDArray]
+WeightFunc = typing.Callable[[ArrayLike1D], npt.NDArray]
 
 
 def distance_weight(z: npt.ArrayLike, cosmo: Cosmology) -> npt.NDArray:
@@ -76,7 +76,7 @@ def density_weight(z: npt.ArrayLike, cosmo: Cosmology) -> npt.NDArray:
     return cosmo.rho_m_z(z) * cosmo.xm(z) ** 2 / cosmo.ef(z)
 
 
-class RadialWindow(NamedTuple):
+class RadialWindow(typing.NamedTuple):
     """
     A radial window, defined by a window function.
 
@@ -120,8 +120,8 @@ class RadialWindow(NamedTuple):
 
     """
 
-    za: Sequence[float]
-    wa: Sequence[float]
+    za: collections.abc.Sequence[float]
+    wa: collections.abc.Sequence[float]
     zeff: float
 
 
@@ -338,7 +338,7 @@ def restrict(
 def partition(
     z: npt.ArrayLike,
     fz: npt.ArrayLike,
-    shells: Sequence[RadialWindow],
+    shells: collections.abc.Sequence[RadialWindow],
     *,
     method: str = "nnls",
 ) -> npt.ArrayLike:
@@ -449,7 +449,7 @@ def partition(
 def partition_lstsq(
     z: npt.ArrayLike,
     fz: npt.ArrayLike,
-    shells: Sequence[RadialWindow],
+    shells: collections.abc.Sequence[RadialWindow],
     *,
     sumtol: float = 0.01,
 ) -> npt.ArrayLike:
@@ -495,7 +495,7 @@ def partition_lstsq(
 def partition_nnls(
     z: npt.ArrayLike,
     fz: npt.ArrayLike,
-    shells: Sequence[RadialWindow],
+    shells: collections.abc.Sequence[RadialWindow],
     *,
     sumtol: float = 0.01,
 ) -> npt.ArrayLike:
@@ -556,7 +556,7 @@ def partition_nnls(
 def partition_restrict(
     z: npt.ArrayLike,
     fz: npt.ArrayLike,
-    shells: Sequence[RadialWindow],
+    shells: collections.abc.Sequence[RadialWindow],
 ) -> npt.ArrayLike:
     """Partition by restriction and integration."""
     part = np.empty((len(shells),) + np.shape(fz)[:-1])
@@ -594,7 +594,7 @@ def distance_grid(cosmo, zmin, zmax, *, dx=None, num=None):
 def combine(
     z: npt.ArrayLike,
     weights: npt.ArrayLike,
-    shells: Sequence[RadialWindow],
+    shells: collections.abc.Sequence[RadialWindow],
 ) -> npt.ArrayLike:
     r"""
     Evaluate a linear combination of window functions.

pyproject.toml

@@ -69,7 +69,7 @@ Issues = "https://github.com/glass-dev/glass/issues"
 
 [tool.coverage]
 report = {exclude_also = [
-    "if TYPE_CHECKING:",
+    "if typing.TYPE_CHECKING:",
 ], omit = [
     "glass/_version.py",
 ], skip_covered = true, sort = "cover"}