fix some tests + minor tweaks

Author: benbovy

xarray/core/indexing.py

@@ -127,17 +127,18 @@ def __init__(self, query_results: List[QueryResult]):
                 + "Suggestion: use a multi-index for each of those dimension(s)."
             )
 
-        self.dim_indexers = {
-            k: v for res in query_results for k, v in res.dim_indexers.items()
-        }
-        self.indexes = {k: v for res in query_results for k, v in res.indexes.items()}
-        self.index_vars = {
-            k: v for res in query_results for k, v in res.index_vars.items()
-        }
-        self.drop_coords = [c for res in query_results for c in res.drop_coords]
-        self.rename_dims = {
-            k: v for res in query_results for k, v in res.rename_dims.items()
-        }
+        self.dim_indexers = {}
+        self.indexes = {}
+        self.index_vars = {}
+        self.drop_coords = []
+        self.rename_dims = {}
+
+        for res in query_results:
+            self.dim_indexers.update(res.dim_indexers)
+            self.indexes.update(res.indexes)
+            self.index_vars.update(res.index_vars)
+            self.drop_coords += res.drop_coords
+            self.rename_dims.update(res.rename_dims)
 
     def to_tuple(self):
         return (
@@ -191,15 +192,14 @@ def group_indexers_by_index(
             unique_indexes[index_id] = index
             label = maybe_cast_to_coords_dtype(label, coord.dtype)  # type: ignore
             grouped_indexers[index_id][key] = label
-        elif coord is not None:
+        elif key in obj.coords:
             raise KeyError(f"no index found for coordinate {key}")
         elif key not in obj.dims:
             raise KeyError(f"{key} is not a valid dimension or coordinate")
         elif len(query_kwargs):
             raise ValueError(
-                "cannot supply selection options "
-                "when the indexed dimension does not have "
-                "an associated coordinate."
+                f"cannot supply selection options {query_kwargs!r} for dimension {key!r}"
+                "that has no asssociated coordinate or index"
             )
         else:
             # key is a dimension without coordinate
@@ -212,14 +212,19 @@ def group_indexers_by_index(
 def remap_label_indexers(
     obj: Union["DataArray", "Dataset"],
     indexers: Mapping[Hashable, Any],
-    **query_kwargs,
-    # query_kwargs: Mapping[str, Any],
-    # **indexers_kwargs,
+    method=None,
+    tolerance=None,
 ) -> MergedQueryResults:
     """Execute index queries from a DataArray / Dataset and label-based indexers
     and return the (merged) query results.
 
     """
+    # TODO benbovy - flexible indexes: remove when custom index options are available
+    if method is None and tolerance is None:
+        query_kwargs = {}
+    else:
+        query_kwargs = {"method": method, "tolerance": tolerance}
+
     # indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "map_index_queries")
     indexes, grouped_indexers = group_indexers_by_index(obj, indexers, query_kwargs)
 

xarray/tests/test_dataset.py

@@ -1680,7 +1680,7 @@ def test_sel_method(self):
 
         with pytest.raises(TypeError, match=r"``method``"):
             # this should not pass silently
-            data.sel(method=data)
+            data.sel(dim2=1, method=data)
 
         # cannot pass method if there is no associated coordinate
         with pytest.raises(ValueError, match=r"cannot supply"):

xarray/tests/test_indexes.py

@@ -86,11 +86,11 @@ def test_query_datetime(self):
             pd.to_datetime(["2000-01-01", "2001-01-01", "2002-01-01"]), "x"
         )
         actual = index.query({"x": "2001-01-01"})
-        expected = ({"x": 1}, None)
-        assert actual == expected
+        expected_dim_indexers = {"x": 1}
+        assert actual.dim_indexers == expected_dim_indexers
 
         actual = index.query({"x": index.to_pandas_index().to_numpy()[1]})
-        assert actual == expected
+        assert actual.dim_indexers == expected_dim_indexers
 
     def test_query_unsorted_datetime_index_raises(self):
         index = PandasIndex(pd.to_datetime(["2001", "2000", "2002"]), "x")
@@ -191,11 +191,11 @@ def test_query(self):
         index = PandasMultiIndex(
             pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=("one", "two")), "x"
         )
+
         # test tuples inside slice are considered as scalar indexer values
-        assert index.query({"x": slice(("a", 1), ("b", 2))}) == (
-            {"x": slice(0, 4)},
-            None,
-        )
+        actual = index.query({"x": slice(("a", 1), ("b", 2))})
+        expected_dim_indexers = {"x": slice(0, 4)}
+        assert actual.dim_indexers == expected_dim_indexers
 
         with pytest.raises(KeyError, match=r"not all values found"):
             index.query({"x": [0]})

xarray/tests/test_indexing.py

@@ -64,21 +64,23 @@ def test_group_indexers_by_index(self):
         data.coords["y2"] = ("y", [2.0, 3.0])
 
         indexes, grouped_indexers = indexing.group_indexers_by_index(
-            data, {"z": 0, "one": "a", "two": 1, "y": 0}
+            data, {"z": 0, "one": "a", "two": 1, "y": 0}, {}
         )
-        assert indexes == {"x": data.xindexes["x"], "y": data.xindexes["y"]}
-        assert grouped_indexers == {
-            "x": {"one": "a", "two": 1},
-            "y": {"y": 0},
-            None: {"z": 0},
-        }
+        for k in indexes:
+            if indexes[k].equals(data.xindexes["x"]):
+                assert grouped_indexers[k] == {"one": "a", "two": 1}
+            elif indexes[k].equals(data.xindexes["y"]):
+                assert grouped_indexers[k] == {"y": 0}
+        assert grouped_indexers[None] == {"z": 0}
+        grouped_indexers.pop(None)
+        assert indexes.keys() == grouped_indexers.keys()
 
         with pytest.raises(KeyError, match=r"no index found for coordinate y2"):
-            indexing.group_indexers_by_index(data, {"y2": 2.0})
+            indexing.group_indexers_by_index(data, {"y2": 2.0}, {})
         with pytest.raises(KeyError, match=r"w is not a valid dimension or coordinate"):
-            indexing.group_indexers_by_index(data, {"w": "a"})
+            indexing.group_indexers_by_index(data, {"w": "a"}, {})
         with pytest.raises(ValueError, match=r"cannot supply.*"):
-            indexing.group_indexers_by_index(data, {"z": 1}, method="nearest")
+            indexing.group_indexers_by_index(data, {"z": 1}, {"method": "nearest"})
 
     def test_remap_label_indexers(self):
         def test_indexer(
@@ -88,6 +90,7 @@ def test_indexer(
             expected_idx=None,
             expected_vars=None,
             expected_drop=None,
+            expected_rename_dims=None,
         ):
             if expected_vars is None:
                 expected_vars = {}
@@ -97,22 +100,23 @@ def test_indexer(
                 expected_idx = {k: expected_idx for k in expected_vars}
             if expected_drop is None:
                 expected_drop = []
+            if expected_rename_dims is None:
+                expected_rename_dims = {}
 
-            pos, new_idx, new_vars, drop_vars = indexing.remap_label_indexers(
-                data, {"x": x}
-            )
+            results = indexing.remap_label_indexers(data, {"x": x})
 
-            assert_array_equal(pos.get("x"), expected_pos)
+            assert_array_equal(results.dim_indexers.get("x"), expected_pos)
 
-            assert new_idx.keys() == expected_idx.keys()
-            for k in new_idx:
-                assert new_idx[k].equals(expected_idx[k])
+            assert results.indexes.keys() == expected_idx.keys()
+            for k in results.indexes:
+                assert results.indexes[k].equals(expected_idx[k])
 
-            assert new_vars.keys() == expected_vars.keys()
-            for k in new_vars:
-                assert_array_equal(new_vars[k], expected_vars[k])
+            assert results.index_vars.keys() == expected_vars.keys()
+            for k in results.index_vars:
+                assert_array_equal(results.index_vars[k], expected_vars[k])
 
-            assert drop_vars == expected_drop
+            assert set(results.drop_coords) == set(expected_drop)
+            assert results.rename_dims == expected_rename_dims
 
         data = Dataset({"x": ("x", [1, 2, 3])})
         mindex = pd.MultiIndex.from_product(
@@ -130,6 +134,7 @@ def test_indexer(
             [True, True, False, False, False, False, False, False],
             *PandasIndex.from_pandas_index(pd.Index([-1, -2]), "three"),
             ["x", "one", "two"],
+            {"x": "three"},
         )
         test_indexer(
             mdata,
@@ -161,13 +166,15 @@ def test_indexer(
             [True, True, False, False, False, False, False, False],
             *PandasIndex.from_pandas_index(pd.Index([-1, -2]), "three"),
             ["x", "one", "two"],
+            {"x": "three"},
         )
         test_indexer(
             mdata,
             {"one": "a", "three": -1},
             [True, False, True, False, False, False, False, False],
             *PandasIndex.from_pandas_index(pd.Index([1, 2]), "two"),
             ["x", "one", "three"],
+            {"x": "two"},
         )
         test_indexer(
             mdata,