Fix type annotations for get_axis_num (GH 9822) (#9827)

* Fix type annotations for `get_axis_num` (GH 9822)

Explicitly annotate that a single `str` argument leads to an `int`
return, overriding the match against Iterable.

* Suppress mypy errors related to get_axis_num

* Add PR number to whats-new.rst

Author: bmerry

doc/whats-new.rst

@@ -36,6 +36,8 @@ Deprecations
 
 Bug fixes
 ~~~~~~~~~
+- Fix type annotations for ``get_axis_num``. (:issue:`9822`, :pull:`9827`).
+  By `Bruce Merry <https://github.com/bmerry>`_.
 - Fix unintended load on datasets when calling :py:meth:`DataArray.plot.scatter` (:pull:`9818`).
   By `Jimmy Westling <https://github.com/illviljan>`_.
 

xarray/core/common.py

@@ -213,6 +213,9 @@ def __iter__(self: Any) -> Iterator[Any]:
             raise TypeError("iteration over a 0-d array")
         return self._iter()
 
+    @overload
+    def get_axis_num(self, dim: str) -> int: ...  # type: ignore [overload-overlap]
+
     @overload
     def get_axis_num(self, dim: Iterable[Hashable]) -> tuple[int, ...]: ...
 

xarray/namedarray/core.py

@@ -669,6 +669,9 @@ 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: str) -> int: ...  # type: ignore [overload-overlap]
+
     @overload
     def get_axis_num(self, dim: Iterable[Hashable]) -> tuple[int, ...]: ...