@@ -59,27 +59,34 @@ def test_consolidate_slices() -> None:
59
59
_consolidate_slices ([slice (3 ), 4 ]) # type: ignore[list-item]
60
60
61
61
62
- def test_groupby_dims_property (dataset ) -> None :
63
- assert dataset .groupby ("x" ).dims == dataset .isel (x = 1 ).dims
64
- assert dataset .groupby ("y" ).dims == dataset .isel (y = 1 ).dims
62
+ def test_groupby_dims_property (dataset , recwarn ) -> None :
63
+ # dims is sensitive to squeeze, always warn
64
+ with pytest .warns (UserWarning , match = "The `squeeze` kwarg" ):
65
+ assert dataset .groupby ("x" ).dims == dataset .isel (x = 1 ).dims
66
+ assert dataset .groupby ("y" ).dims == dataset .isel (y = 1 ).dims
65
67
68
+ # when squeeze=False, no warning should be raised
66
69
assert dataset .groupby ("x" , squeeze = False ).dims == dataset .isel (x = slice (1 , 2 )).dims
67
70
assert dataset .groupby ("y" , squeeze = False ).dims == dataset .isel (y = slice (1 , 2 )).dims
71
+ assert len (recwarn ) == 0
68
72
69
73
stacked = dataset .stack ({"xy" : ("x" , "y" )})
70
74
assert stacked .groupby ("xy" , squeeze = False ).dims == stacked .isel (xy = [0 ]).dims
75
+ assert len (recwarn ) == 0
71
76
72
77
73
78
def test_multi_index_groupby_map (dataset ) -> None :
74
79
# regression test for GH873
75
80
ds = dataset .isel (z = 1 , drop = True )[["foo" ]]
76
81
expected = 2 * ds
77
- actual = (
78
- ds .stack (space = ["x" , "y" ])
79
- .groupby ("space" )
80
- .map (lambda x : 2 * x )
81
- .unstack ("space" )
82
- )
82
+ # The function in `map` may be sensitive to squeeze, always warn
83
+ with pytest .warns (UserWarning , match = "The `squeeze` kwarg" ):
84
+ actual = (
85
+ ds .stack (space = ["x" , "y" ])
86
+ .groupby ("space" )
87
+ .map (lambda x : 2 * x )
88
+ .unstack ("space" )
89
+ )
83
90
assert_equal (expected , actual )
84
91
85
92
@@ -202,7 +209,9 @@ def func(arg1, arg2, arg3=0):
202
209
203
210
dataset = xr .Dataset ({"foo" : ("x" , [1 , 1 , 1 ])}, {"x" : [1 , 2 , 3 ]})
204
211
expected = xr .Dataset ({"foo" : ("x" , [3 , 3 , 3 ])}, {"x" : [1 , 2 , 3 ]})
205
- actual = dataset .groupby ("x" ).map (func , args = (1 ,), arg3 = 1 )
212
+ # The function in `map` may be sensitive to squeeze, always warn
213
+ with pytest .warns (UserWarning , match = "The `squeeze` kwarg" ):
214
+ actual = dataset .groupby ("x" ).map (func , args = (1 ,), arg3 = 1 )
206
215
assert_identical (expected , actual )
207
216
208
217
@@ -887,7 +896,7 @@ def test_groupby_bins_cut_kwargs(use_flox: bool) -> None:
887
896
888
897
with xr .set_options (use_flox = use_flox ):
889
898
actual = da .groupby_bins (
890
- "x" , bins = x_bins , include_lowest = True , right = False
899
+ "x" , bins = x_bins , include_lowest = True , right = False , squeeze = False
891
900
).mean ()
892
901
expected = xr .DataArray (
893
902
np .array ([[1.0 , 2.0 ], [5.0 , 6.0 ], [9.0 , 10.0 ]]),
@@ -1135,8 +1144,8 @@ def test_groupby_properties(self):
1135
1144
"by, use_da" , [("x" , False ), ("y" , False ), ("y" , True ), ("abc" , False )]
1136
1145
)
1137
1146
@pytest .mark .parametrize ("shortcut" , [True , False ])
1138
- @pytest .mark .parametrize ("squeeze" , [True , False ])
1139
- def test_groupby_map_identity (self , by , use_da , shortcut , squeeze ) -> None :
1147
+ @pytest .mark .parametrize ("squeeze" , [None ])
1148
+ def test_groupby_map_identity (self , by , use_da , shortcut , squeeze , recwarn ) -> None :
1140
1149
expected = self .da
1141
1150
if use_da :
1142
1151
by = expected .coords [by ]
@@ -1148,6 +1157,10 @@ def identity(x):
1148
1157
actual = grouped .map (identity , shortcut = shortcut )
1149
1158
assert_identical (expected , actual )
1150
1159
1160
+ # abc is not a dim coordinate so no warnings expected!
1161
+ if (by .name if use_da else by ) != "abc" :
1162
+ assert len (recwarn ) == (1 if squeeze in [None , True ] else 0 )
1163
+
1151
1164
def test_groupby_sum (self ):
1152
1165
array = self .da
1153
1166
grouped = array .groupby ("abc" )
@@ -1508,7 +1521,7 @@ def test_groupby_bins_ellipsis(self):
1508
1521
da = xr .DataArray (np .ones ((2 , 3 , 4 )))
1509
1522
bins = [- 1 , 0 , 1 , 2 ]
1510
1523
with xr .set_options (use_flox = False ):
1511
- actual = da .groupby_bins ("dim_0" , bins ).mean (...)
1524
+ actual = da .groupby_bins ("dim_0" , bins , squeeze = False ).mean (...)
1512
1525
with xr .set_options (use_flox = True ):
1513
1526
expected = da .groupby_bins ("dim_0" , bins ).mean (...)
1514
1527
assert_allclose (actual , expected )
0 commit comments