diff --git a/lib/iris/coords.py b/lib/iris/coords.py index 8afe9dad41..97a50e5608 100644 --- a/lib/iris/coords.py +++ b/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,8 +2106,8 @@ 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. """ @@ -2109,13 +2115,15 @@ def collapsed(self, dims_to_collapse=None): # 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())