Restore old MultiIndex dropping behaviour (pydata#6592)

Co-authored-by: Benoit Bovy <bbovy@gfz-potsdam.de>
HertugHelms · May 11, 2022 · 11041bd · 11041bd
1 parent f6d0f84
commit 11041bd
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 9 deletions.
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -4551,6 +4551,23 @@ def drop_vars(
         if errors == "raise":
             self._assert_all_in_dataset(names)
 
+        # GH6505
+        other_names = set()
+        for var in names:
+            maybe_midx = self._indexes.get(var, None)
+            if isinstance(maybe_midx, PandasMultiIndex):
+                idx_coord_names = set(maybe_midx.index.names + [maybe_midx.dim])
+                idx_other_names = idx_coord_names - set(names)
+                other_names.update(idx_other_names)
+        if other_names:
+            names |= set(other_names)
+            warnings.warn(
+                f"Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. "
+                f"Please also drop the following variables: {other_names!r} to avoid an error in the future.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         assert_no_index_corrupted(self.xindexes, names)
 
         variables = {k: v for k, v in self._variables.items() if k not in names}

diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
@@ -2360,10 +2360,11 @@ def test_drop_coordinates(self):
         assert_identical(actual, renamed)
 
     def test_drop_multiindex_level(self):
-        with pytest.raises(
-            ValueError, match=r"cannot remove coordinate.*corrupt.*index "
-        ):
-            self.mda.drop_vars("level_1")
+        # GH6505
+        expected = self.mda.drop_vars(["x", "level_1", "level_2"])
+        with pytest.warns(DeprecationWarning):
+            actual = self.mda.drop_vars("level_1")
+        assert_identical(expected, actual)
 
     def test_drop_all_multiindex_levels(self):
         dim_levels = ["x", "level_1", "level_2"]

diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
@@ -2453,11 +2453,10 @@ def test_drop_variables(self):
 
     def test_drop_multiindex_level(self):
         data = create_test_multiindex()
-
-        with pytest.raises(
-            ValueError, match=r"cannot remove coordinate.*corrupt.*index "
-        ):
-            data.drop_vars("level_1")
+        expected = data.drop_vars(["x", "level_1", "level_2"])
+        with pytest.warns(DeprecationWarning):
+            actual = data.drop_vars("level_1")
+        assert_identical(expected, actual)
 
     def test_drop_index_labels(self):
         data = Dataset({"A": (["x", "y"], np.random.randn(2, 3)), "x": ["a", "b"]})