Add overload to get_axis_num

pydata · Dec 12, 2023 · 6231ae2 · 6231ae2
1 parent 0bf38c2
commit 6231ae2
Showing 1 changed file with 8 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.