Refactor out test Index classes

Author: dcherian

xarray/tests/indexes.py

@@ -0,0 +1,53 @@
+from xarray.core.indexes import Index, PandasIndex
+
+
+class ScalarIndex(Index):
+    def __init__(self, value: int):
+        self.value = value
+
+    @classmethod
+    def from_variables(cls, variables, *, options):
+        var = next(iter(variables.values()))
+        return cls(int(var.values))
+
+    def equals(self, other, *, exclude=None):
+        return isinstance(other, ScalarIndex) and other.value == self.value
+
+
+class XYIndex(Index):
+    def __init__(self, x: PandasIndex, y: PandasIndex):
+        self.x: PandasIndex = x
+        self.y: PandasIndex = y
+
+    @classmethod
+    def from_variables(cls, variables, *, options):
+        return cls(
+            x=PandasIndex.from_variables({"x": variables["x"]}, options=options),
+            y=PandasIndex.from_variables({"y": variables["y"]}, options=options),
+        )
+
+    def equals(self, other, exclude=None):
+        x_eq = True if self.x.dim in exclude else self.x.equals(other.x)
+        y_eq = True if self.y.dim in exclude else self.y.equals(other.y)
+        return x_eq and y_eq
+
+
+class MultiCoordIndex(Index):
+    def __init__(self, idx1, idx2):
+        self.idx1 = idx1
+        self.idx2 = idx2
+
+    @classmethod
+    def from_variables(cls, variables, *, options=None):
+        idx1 = PandasIndex.from_variables({"x": variables["x"]}, options=options)
+        idx2 = PandasIndex.from_variables({"y": variables["y"]}, options=options)
+
+        return cls(idx1, idx2)
+
+    def create_variables(self, variables=None):
+        return {**self.idx1.create_variables(), **self.idx2.create_variables()}
+
+    def isel(self, indexers):
+        idx1 = self.idx1.isel({"x": indexers.get("x", slice(None))})
+        idx2 = self.idx2.isel({"y": indexers.get("y", slice(None))})
+        return MultiCoordIndex(idx1, idx2)

xarray/tests/test_dataset.py

@@ -69,6 +69,7 @@
     requires_sparse,
     source_ndarray,
 )
+from xarray.tests.indexes import MultiCoordIndex, ScalarIndex, XYIndex
 
 try:
     from pandas.errors import UndefinedVariableError
@@ -1599,30 +1600,6 @@ def test_isel_multicoord_index(self) -> None:
         # regression test https://github.com/pydata/xarray/issues/10063
         # isel on a multi-coordinate index should return a unique index associated
         # to each coordinate
-        class MultiCoordIndex(xr.Index):
-            def __init__(self, idx1, idx2):
-                self.idx1 = idx1
-                self.idx2 = idx2
-
-            @classmethod
-            def from_variables(cls, variables, *, options=None):
-                idx1 = PandasIndex.from_variables(
-                    {"x": variables["x"]}, options=options
-                )
-                idx2 = PandasIndex.from_variables(
-                    {"y": variables["y"]}, options=options
-                )
-
-                return cls(idx1, idx2)
-
-            def create_variables(self, variables=None):
-                return {**self.idx1.create_variables(), **self.idx2.create_variables()}
-
-            def isel(self, indexers):
-                idx1 = self.idx1.isel({"x": indexers.get("x", slice(None))})
-                idx2 = self.idx2.isel({"y": indexers.get("y", slice(None))})
-                return MultiCoordIndex(idx1, idx2)
-
         coords = xr.Coordinates(coords={"x": [0, 1], "y": [1, 2]}, indexes={})
         ds = xr.Dataset(coords=coords).set_xindex(["x", "y"], MultiCoordIndex)
 
@@ -2639,18 +2616,6 @@ def test_align_index_var_attrs(self, join) -> None:
     def test_align_scalar_index(self) -> None:
         # ensure that indexes associated with scalar coordinates are not ignored
         # during alignment
-        class ScalarIndex(Index):
-            def __init__(self, value: int):
-                self.value = value
-
-            @classmethod
-            def from_variables(cls, variables, *, options):
-                var = next(iter(variables.values()))
-                return cls(int(var.values))
-
-            def equals(self, other, *, exclude=None):
-                return isinstance(other, ScalarIndex) and other.value == self.value
-
         ds1 = Dataset(coords={"x": 0}).set_xindex("x", ScalarIndex)
         ds2 = Dataset(coords={"x": 0}).set_xindex("x", ScalarIndex)
 
@@ -2664,27 +2629,6 @@ def equals(self, other, *, exclude=None):
             xr.align(ds1, ds3, join="exact")
 
     def test_align_multi_dim_index_exclude_dims(self) -> None:
-        class XYIndex(Index):
-            def __init__(self, x: PandasIndex, y: PandasIndex):
-                self.x: PandasIndex = x
-                self.y: PandasIndex = y
-
-            @classmethod
-            def from_variables(cls, variables, *, options):
-                return cls(
-                    x=PandasIndex.from_variables(
-                        {"x": variables["x"]}, options=options
-                    ),
-                    y=PandasIndex.from_variables(
-                        {"y": variables["y"]}, options=options
-                    ),
-                )
-
-            def equals(self, other, exclude=None):
-                x_eq = True if self.x.dim in exclude else self.x.equals(other.x)
-                y_eq = True if self.y.dim in exclude else self.y.equals(other.y)
-                return x_eq and y_eq
-
         ds1 = (
             Dataset(coords={"x": [1, 2], "y": [3, 4]})
             .drop_indexes(["x", "y"])