Add overloads to get_axis_num (#8547)

Co-authored-by: Anderson Banihirwe <[email protected]>
pydata · Jan 30, 2024 · 81f38f3 · 81f38f3
1 parent b0b5b2f
commit 81f38f3
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/xarray/core/common.py b/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.
 

diff --git a/xarray/namedarray/core.py b/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.