From 29e312f3ceea501ff6a886244d67e758329d96e1 Mon Sep 17 00:00:00 2001
From: Ruth Comer <ruth.comer@metoffice.gov.uk>
Date: Wed, 5 Jan 2022 13:24:33 +0000
Subject: [PATCH] type hint improvements

---
 lib/iris/coords.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/lib/iris/coords.py b/lib/iris/coords.py
index ffd3df5504..08afbada68 100644
--- a/lib/iris/coords.py
+++ b/lib/iris/coords.py
@@ -14,7 +14,7 @@
 import copy
 from itertools import chain, 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):
     """
@@ -250,7 +255,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.
@@ -1447,7 +1452,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.
@@ -1565,7 +1570,7 @@ 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.
@@ -1573,7 +1578,7 @@ def core_points(self):
         """
         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.
@@ -1935,9 +1940,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.
@@ -1956,8 +1959,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())