diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 68c20c4f51b..820fcd48bd3 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -118,8 +118,10 @@ def test_apply_identity() -> None: assert_identical(variable, apply_identity(variable)) assert_identical(data_array, apply_identity(data_array)) assert_identical(data_array, apply_identity(data_array.groupby("x"))) + assert_identical(data_array, apply_identity(data_array.groupby("x", squeeze=False))) assert_identical(dataset, apply_identity(dataset)) assert_identical(dataset, apply_identity(dataset.groupby("x"))) + assert_identical(dataset, apply_identity(dataset.groupby("x", squeeze=False))) def add(a, b): @@ -519,8 +521,10 @@ def func(x): assert_identical(stacked_variable, stack_negative(variable)) assert_identical(stacked_data_array, stack_negative(data_array)) assert_identical(stacked_dataset, stack_negative(dataset)) - assert_identical(stacked_data_array, stack_negative(data_array.groupby("x"))) - assert_identical(stacked_dataset, stack_negative(dataset.groupby("x"))) + with pytest.warns(UserWarning, match="The `squeeze` kwarg"): + assert_identical(stacked_data_array, stack_negative(data_array.groupby("x"))) + with pytest.warns(UserWarning, match="The `squeeze` kwarg"): + assert_identical(stacked_dataset, stack_negative(dataset.groupby("x"))) def original_and_stack_negative(obj): def func(x): @@ -547,11 +551,13 @@ def func(x): assert_identical(dataset, out0) assert_identical(stacked_dataset, out1) - out0, out1 = original_and_stack_negative(data_array.groupby("x")) + with pytest.warns(UserWarning, match="The `squeeze` kwarg"): + out0, out1 = original_and_stack_negative(data_array.groupby("x")) assert_identical(data_array, out0) assert_identical(stacked_data_array, out1) - out0, out1 = original_and_stack_negative(dataset.groupby("x")) + with pytest.warns(UserWarning, match="The `squeeze` kwarg"): + out0, out1 = original_and_stack_negative(dataset.groupby("x")) assert_identical(dataset, out0) assert_identical(stacked_dataset, out1) diff --git a/xarray/tests/test_concat.py b/xarray/tests/test_concat.py index d1fc085bf0f..0cf4cc03a09 100644 --- a/xarray/tests/test_concat.py +++ b/xarray/tests/test_concat.py @@ -494,7 +494,7 @@ def test_concat_merge_variables_present_in_some_datasets(self, data) -> None: def test_concat_2(self, data) -> None: dim = "dim2" - datasets = [g for _, g in data.groupby(dim, squeeze=True)] + datasets = [g.squeeze(dim) for _, g in data.groupby(dim, squeeze=False)] concat_over = [k for k, v in data.coords.items() if dim in v.dims and k != dim] actual = concat(datasets, data[dim], coords=concat_over) assert_identical(data, self.rectify_dim_order(data, actual)) @@ -505,7 +505,7 @@ def test_concat_coords_kwarg(self, data, dim, coords) -> None: data = data.copy(deep=True) # make sure the coords argument behaves as expected data.coords["extra"] = ("dim4", np.arange(3)) - datasets = [g for _, g in data.groupby(dim, squeeze=True)] + datasets = [g.squeeze() for _, g in data.groupby(dim, squeeze=False)] actual = concat(datasets, data[dim], coords=coords) if coords == "all": @@ -1000,7 +1000,7 @@ def test_concat(self) -> None: actual = concat([foo, bar], "w") assert_equal(expected, actual) # from iteration: - grouped = [g for _, g in foo.groupby("x")] + grouped = [g.squeeze() for _, g in foo.groupby("x", squeeze=False)] stacked = concat(grouped, ds["x"]) assert_identical(foo, stacked) # with an index as the 'dim' argument diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 35259094f54..cd0df4b93b1 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -1475,10 +1475,9 @@ def test_groupby_restore_coord_dims(self): ("a", ("a", "y")), ("b", ("x", "b")), ]: - with pytest.warns(UserWarning, match="The `squeeze` kwarg"): - result = array.groupby(by, restore_coord_dims=True).map( - lambda x: x.squeeze() - )["c"] + result = array.groupby(by, squeeze=False, restore_coord_dims=True).map( + lambda x: x.squeeze() + )["c"] assert result.dims == expected_dims def test_groupby_first_and_last(self):