Some type annotations for coords.py

lib/iris/coords.py

@@ -11,12 +11,14 @@
 from functools import lru_cache
 from itertools import zip_longest
 import operator
+from typing import Iterable, Optional, Union
 import warnings
 import zlib
 
 import dask.array as da
 import numpy as np
 import numpy.ma as ma
+import numpy.typing as npt
 
 from iris._data_manager import DataManager
 import iris._lazy_data as _lazy
@@ -32,11 +34,15 @@
 import iris.exceptions
 import iris.time
 import iris.util
-import iris.warnings
 
 #: The default value for ignore_axis which controls guess_coord_axis' behaviour
 DEFAULT_IGNORE_AXIS = False
 
+# 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):
     """Superclass for dimensional metadata."""
@@ -238,7 +244,7 @@ def _lazy_values(self):
         """Return a lazy array representing the dimensional metadata values."""
         return self._values_dm.lazy_data()
 
-    def _core_values(self):
+    def _core_values(self) -> RealOrLazyData:
         """Value array of this dimensional metadata which may be a NumPy array or a dask array."""
         result = self._values_dm.core_data()
         if not _lazy.is_lazy_data(result):
@@ -771,7 +777,7 @@ def dtype(self):
         return self._values_dm.dtype
 
     @property
-    def ndim(self):
+    def ndim(self) -> int:
         """Return the number of dimensions of the current dimensional metadata object."""
         return self._values_dm.ndim
 
@@ -1585,7 +1591,7 @@ def points(self, points):
         self._values = points
 
     @property
-    def bounds(self):
+    def bounds(self) -> Optional[RealData]:
         """Coordinate bounds values.
 
         The coordinate bounds values, as a NumPy array,
@@ -1717,11 +1723,11 @@ def lazy_bounds(self):
             lazy_bounds = self._bounds_dm.lazy_data()
         return lazy_bounds
 
-    def core_points(self):
+    def core_points(self) -> RealOrLazyData:
         """Core 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):
+    def core_bounds(self) -> Optional[RealOrLazyData]:
         """Core bounds. The points array at the core of this coord, which may be a NumPy array or a dask array."""
         result = None
         if self.has_bounds():
@@ -2100,22 +2106,24 @@ def cell(self, index):
 
         return Cell(point, bound)
 
-    def collapsed(self, dims_to_collapse=None):
-        """Return a copy of this coordinate, which has been collapsed along the specified dimensions.
+    def collapsed(self, dims_to_collapse: Optional[Dims] = None) -> "Coord":
+        """Returns a copy of this coordinate, which has been collapsed along the specified dimensions.
 
         Replaces the points & bounds with a simple bounded region.
         """
         # Ensure dims_to_collapse is a tuple to be able to pass
         # through to numpy
         if isinstance(dims_to_collapse, (int, np.integer)):
             dims_to_collapse = (dims_to_collapse,)
-        if isinstance(dims_to_collapse, list):
+        if isinstance(dims_to_collapse, Iterable):
             dims_to_collapse = tuple(dims_to_collapse)
 
         if np.issubdtype(self.dtype, np.str_):
             # Collapse the coordinate by serializing the points and
             # bounds as strings.
-            def serialize(x, axis):
+            def serialize(
+                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())