type hint improvements

Author: rcomer

lib/iris/coords.py

@@ -15,7 +15,7 @@
 from functools import lru_cache
 from itertools import zip_longest
 import operator
-from typing import Sequence, Union
+from typing import Iterable, Optional, Union
 import warnings
 import zlib
 
@@ -39,6 +39,11 @@
 import iris.time
 import iris.util
 
+# Define some typing aliases.
+Dims = Union[int, Iterable[int]]
+RealData = Union[np.ndarray, ma.MaskedArray]
+RealOrLazyData = Union[RealData, da.Array]
+
 
 class _DimensionalMetadata(CFVariableMixin, metaclass=ABCMeta):
     """
@@ -254,7 +259,7 @@ def _lazy_values(self):
         """
         return self._values_dm.lazy_data()
 
-    def _core_values(self) -> Union[npt.NDArray, "da.Array"]:
+    def _core_values(self) -> RealOrLazyData:
         """
         The values array of this dimensional metadata which may be a NumPy
         array or a dask array.
@@ -1659,7 +1664,7 @@ def points(self, points):
         self._values = points
 
     @property
-    def bounds(self) -> npt.NDArray:
+    def bounds(self) -> RealData:
         """
         The coordinate bounds values, as a NumPy array,
         or None if no bound values are defined.
@@ -1777,15 +1782,15 @@ def lazy_bounds(self):
             lazy_bounds = self._bounds_dm.lazy_data()
         return lazy_bounds
 
-    def core_points(self):
+    def core_points(self) -> RealOrLazyData:
         """
         The points array at the core of this coord, which may be a NumPy array
         or a dask array.
 
         """
         return super()._core_values()
 
-    def core_bounds(self) -> Union[npt.NDArray, "da.Array"]:
+    def core_bounds(self) -> RealOrLazyData:
         """
         The points array at the core of this coord, which may be a NumPy array
         or a dask array.
@@ -2184,9 +2189,7 @@ def cell(self, index):
 
         return Cell(point, bound)
 
-    def collapsed(
-        self, dims_to_collapse: Union[int, Sequence[int], None] = None
-    ) -> "Coord":
+    def collapsed(self, dims_to_collapse: Optional[Dims] = None) -> "Coord":
         """
         Returns a copy of this coordinate, which has been collapsed along
         the specified dimensions.
@@ -2205,8 +2208,8 @@ def collapsed(
             # Collapse the coordinate by serializing the points and
             # bounds as strings.
             def serialize(
-                x: npt.NDArray, axis: Union[Sequence[int], None]
-            ) -> Union[npt.NDArray, str]:
+                x: npt.NDArray[np.str_], axis: Optional[Iterable[int]]
+            ) -> Union[npt.NDArray[np.str_], str]:
                 if axis is None:
                     return "|".join(str(i) for i in x.flatten())