From 6231ae29a4d7c18b129b16b0e3c29d44eb17406b Mon Sep 17 00:00:00 2001
From: Illviljan <14371165+Illviljan@users.noreply.github.com>
Date: Tue, 12 Dec 2023 22:15:29 +0100
Subject: [PATCH 1/2] Add overload to get_axis_num

---
 xarray/core/common.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/xarray/core/common.py b/xarray/core/common.py
index 6dff9cc4024..9f3c8eabf7b 100644
--- 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.
 

From 48caa38e85734a2faf29ab7b96e3df9fbc95bcae Mon Sep 17 00:00:00 2001
From: Illviljan <14371165+Illviljan@users.noreply.github.com>
Date: Tue, 12 Dec 2023 22:32:45 +0100
Subject: [PATCH 2/2] fix in namedarray

---
 xarray/namedarray/core.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/xarray/namedarray/core.py b/xarray/namedarray/core.py
index b9ad27b6679..c1564caadc7 100644
--- a/xarray/namedarray/core.py
+++ b/xarray/namedarray/core.py
@@ -642,6 +642,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.