Skip to content

Commit 22ad7fa

Browse files
dcheriandcherian
committed
[WIP]
1 parent c905b74 commit 22ad7fa

File tree

4 files changed

+153
-104
lines changed

4 files changed

+153
-104
lines changed

‎xarray/core/common.py‎

Lines changed: 73 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -814,6 +814,74 @@ def rolling_exp(
814814

815815
return rolling_exp.RollingExp(self, window, window_type)
816816

817+
def _groupby(self, groupby_cls, group, squeeze: bool, restore_coord_dims):
818+
from xarray.core.groupby import UniqueGrouper, _validate_group
819+
820+
# While we don't generally check the type of every arg, passing
821+
# multiple dimensions as multiple arguments is common enough, and the
822+
# consequences hidden enough (strings evaluate as true) to warrant
823+
# checking here.
824+
# A future version could make squeeze kwarg only, but would face
825+
# backward-compat issues.
826+
if not isinstance(squeeze, bool):
827+
raise TypeError(
828+
f"`squeeze` must be True or False, but {squeeze} was supplied"
829+
)
830+
831+
newobj, name = _validate_group(self, group)
832+
833+
grouper = UniqueGrouper()
834+
return groupby_cls(
835+
newobj,
836+
{name: grouper},
837+
squeeze=squeeze,
838+
restore_coord_dims=restore_coord_dims,
839+
)
840+
841+
def _groupby_bins(
842+
self,
843+
groupby_cls,
844+
group: Hashable | DataArray | IndexVariable,
845+
bins: ArrayLike,
846+
right: bool = True,
847+
labels: ArrayLike | None = None,
848+
precision: int = 3,
849+
include_lowest: bool = False,
850+
squeeze: bool = True,
851+
restore_coord_dims: bool = False,
852+
):
853+
from xarray.core.groupby import BinGrouper, _validate_group
854+
855+
# While we don't generally check the type of every arg, passing
856+
# multiple dimensions as multiple arguments is common enough, and the
857+
# consequences hidden enough (strings evaluate as true) to warrant
858+
# checking here.
859+
# A future version could make squeeze kwarg only, but would face
860+
# backward-compat issues.
861+
if not isinstance(squeeze, bool):
862+
raise TypeError(
863+
f"`squeeze` must be True or False, but {squeeze} was supplied"
864+
)
865+
866+
newobj, name = _validate_group(self, group)
867+
868+
grouper = BinGrouper(
869+
bins=bins,
870+
cut_kwargs={
871+
"right": right,
872+
"labels": labels,
873+
"precision": precision,
874+
"include_lowest": include_lowest,
875+
},
876+
)
877+
878+
return groupby_cls(
879+
newobj,
880+
{name: grouper},
881+
squeeze=squeeze,
882+
restore_coord_dims=restore_coord_dims,
883+
)
884+
817885
def _resample(
818886
self,
819887
resample_cls: type[T_Resample],
@@ -1000,12 +1068,13 @@ def _resample(
10001068
if base is not None:
10011069
offset = _convert_base_to_offset(base, freq, index)
10021070

1071+
name = RESAMPLE_DIM
10031072
group = DataArray(
1004-
dim_coord, coords=dim_coord.coords, dims=dim_coord.dims, name=RESAMPLE_DIM
1073+
dim_coord, coords=dim_coord.coords, dims=dim_coord.dims, name=name
10051074
)
1075+
newobj = self.copy().assign_coords({name: group})
10061076

10071077
grouper = TimeResampleGrouper(
1008-
group=group,
10091078
freq=freq,
10101079
closed=closed,
10111080
label=label,
@@ -1015,8 +1084,8 @@ def _resample(
10151084
)
10161085

10171086
return resample_cls(
1018-
self,
1019-
grouper=grouper,
1087+
newobj,
1088+
{name: grouper},
10201089
dim=dim_name,
10211090
resample_dim=RESAMPLE_DIM,
10221091
restore_coord_dims=restore_coord_dims,

‎xarray/core/dataarray.py‎

Lines changed: 13 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -6256,22 +6256,13 @@ def groupby(
62566256
core.groupby.DataArrayGroupBy
62576257
pandas.DataFrame.groupby
62586258
"""
6259-
from xarray.core.groupby import DataArrayGroupBy, UniqueGrouper
6260-
6261-
# While we don't generally check the type of every arg, passing
6262-
# multiple dimensions as multiple arguments is common enough, and the
6263-
# consequences hidden enough (strings evaluate as true) to warrant
6264-
# checking here.
6265-
# A future version could make squeeze kwarg only, but would face
6266-
# backward-compat issues.
6267-
if not isinstance(squeeze, bool):
6268-
raise TypeError(
6269-
f"`squeeze` must be True or False, but {squeeze} was supplied"
6270-
)
6259+
from xarray.core.groupby import DataArrayGroupBy
62716260

6272-
grouper = UniqueGrouper(group)
6273-
return DataArrayGroupBy(
6274-
self, grouper, squeeze=squeeze, restore_coord_dims=restore_coord_dims
6261+
return self._groupby(
6262+
groupby_cls=DataArrayGroupBy,
6263+
group=group,
6264+
squeeze=squeeze,
6265+
restore_coord_dims=restore_coord_dims,
62756266
)
62766267

62776268
def groupby_bins(
@@ -6342,33 +6333,16 @@ def groupby_bins(
63426333
----------
63436334
.. [1] http://pandas.pydata.org/pandas-docs/stable/generated/pandas.cut.html
63446335
"""
6345-
from xarray.core.groupby import BinGrouper, DataArrayGroupBy
6346-
6347-
# While we don't generally check the type of every arg, passing
6348-
# multiple dimensions as multiple arguments is common enough, and the
6349-
# consequences hidden enough (strings evaluate as true) to warrant
6350-
# checking here.
6351-
# A future version could make squeeze kwarg only, but would face
6352-
# backward-compat issues.
6353-
if not isinstance(squeeze, bool):
6354-
raise TypeError(
6355-
f"`squeeze` must be True or False, but {squeeze} was supplied"
6356-
)
6336+
from xarray.core.groupby import DataArrayGroupBy
63576337

6358-
grouper = BinGrouper(
6338+
return self._groupby_bins(
6339+
groupby_cls=DataArrayGroupBy,
63596340
group=group,
63606341
bins=bins,
6361-
cut_kwargs={
6362-
"right": right,
6363-
"labels": labels,
6364-
"precision": precision,
6365-
"include_lowest": include_lowest,
6366-
},
6367-
)
6368-
6369-
return DataArrayGroupBy(
6370-
self,
6371-
grouper,
6342+
right=right,
6343+
labels=labels,
6344+
precision=precision,
6345+
include_lowest=include_lowest,
63726346
squeeze=squeeze,
63736347
restore_coord_dims=restore_coord_dims,
63746348
)

‎xarray/core/dataset.py‎

Lines changed: 15 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -8943,23 +8943,13 @@ def groupby(
89438943
Dataset.resample
89448944
DataArray.resample
89458945
"""
8946-
from xarray.core.groupby import DatasetGroupBy, UniqueGrouper
8947-
8948-
# While we don't generally check the type of every arg, passing
8949-
# multiple dimensions as multiple arguments is common enough, and the
8950-
# consequences hidden enough (strings evaluate as true) to warrant
8951-
# checking here.
8952-
# A future version could make squeeze kwarg only, but would face
8953-
# backward-compat issues.
8954-
if not isinstance(squeeze, bool):
8955-
raise TypeError(
8956-
f"`squeeze` must be True or False, but {squeeze} was supplied"
8957-
)
8946+
from xarray.core.groupby import DatasetGroupBy
89588947

8959-
grouper = UniqueGrouper(group)
8960-
8961-
return DatasetGroupBy(
8962-
self, grouper, squeeze=squeeze, restore_coord_dims=restore_coord_dims
8948+
return self._groupby(
8949+
groupby_cls=DatasetGroupBy,
8950+
group=group,
8951+
squeeze=squeeze,
8952+
restore_coord_dims=restore_coord_dims,
89638953
)
89648954

89658955
def groupby_bins(
@@ -9030,21 +9020,18 @@ def groupby_bins(
90309020
----------
90319021
.. [1] http://pandas.pydata.org/pandas-docs/stable/generated/pandas.cut.html
90329022
"""
9033-
from xarray.core.groupby import BinGrouper, DatasetGroupBy
9023+
from xarray.core.groupby import DatasetGroupBy
90349024

9035-
grouper = BinGrouper(
9025+
return self._groupby_bins(
9026+
groupby_cls=DatasetGroupBy,
90369027
group=group,
90379028
bins=bins,
9038-
cut_kwargs={
9039-
"right": right,
9040-
"labels": labels,
9041-
"precision": precision,
9042-
"include_lowest": include_lowest,
9043-
},
9044-
)
9045-
9046-
return DatasetGroupBy(
9047-
self, grouper, squeeze=squeeze, restore_coord_dims=restore_coord_dims
9029+
right=right,
9030+
labels=labels,
9031+
precision=precision,
9032+
include_lowest=include_lowest,
9033+
squeeze=squeeze,
9034+
restore_coord_dims=restore_coord_dims,
90489035
)
90499036

90509037
def weighted(self, weights: DataArray) -> DatasetWeighted:

0 commit comments

Comments
 (0)