Add overloads to get_axis_num (#8547)

Co-authored-by: Anderson Banihirwe <[email protected]>

Author: Illviljan

xarray/core/common.py

@@ -199,6 +199,14 @@ def __iter__(self: Any) -> Iterator[Any]:
             raise TypeError("iteration over a 0-d array")
         return self._iter()
 
+    @overload
+    def get_axis_num(self, dim: Iterable[Hashable]) -> tuple[int, ...]:
+        ...
+
+    @overload
+    def get_axis_num(self, dim: Hashable) -> int:
+        ...
+
     def get_axis_num(self, dim: Hashable | Iterable[Hashable]) -> int | tuple[int, ...]:
         """Return axis number(s) corresponding to dimension(s) in this array.
 

xarray/namedarray/core.py

@@ -648,6 +648,14 @@ def _dask_finalize(
         data = array_func(results, *args, **kwargs)
         return type(self)(self._dims, data, attrs=self._attrs)
 
+    @overload
+    def get_axis_num(self, dim: Iterable[Hashable]) -> tuple[int, ...]:
+        ...
+
+    @overload
+    def get_axis_num(self, dim: Hashable) -> int:
+        ...
+
     def get_axis_num(self, dim: Hashable | Iterable[Hashable]) -> int | tuple[int, ...]:
         """Return axis number(s) corresponding to dimension(s) in this array.