Function test_groupby_scans

xarray/tests/test_groupby.py:2634–2723

(
    method: Literal["cumsum", "cumprod"],
    grp_idx: list[str],
    dim,
    expected_array: list[float],
    use_flox: bool,
    use_dask: bool,
    use_lazy_group_idx: bool,
)

    ],
)
def test_groupby_scans(
    method: Literal["cumsum", "cumprod"],
    grp_idx: list[str],
    dim,
    expected_array: list[float],
    use_flox: bool,
    use_dask: bool,
    use_lazy_group_idx: bool,
) -> None:
    if use_dask and not has_dask:
        pytest.skip("requires dask")

    if use_flox:
        if not has_flox:
            pytest.skip("requires flox")

        if method == "cumprod":
            pytest.skip(
                "TODO: Groupby with cumprod is currently not supported with flox"
            )
        if dim == ...:
            pytest.skip(
                "TODO: Scans are only supported along a single dimension in flox."
            )
        elif dim == "test":
            pytest.skip(
                "TODO: group_idx along time dim and axis along test dim not currently supported with flox."
            )
    elif use_lazy_group_idx:
        pytest.skip("Lazy group_idx is not supported without flox.")

    # Test Dataset groupby:
    ds = xr.Dataset(
        {
            "foo": (
                ("test", "time"),
                [[7, 2, 0, 1, 2, np.nan], [1, 1, 1, 1, 1, 1], [2, 2, 2, 2, 2, 2]],
            )
        },
        coords={
            "time": [0, 1 / 6, 2 / 6, 3 / 6, 4 / 6, 5 / 6],
            "test": ["a", "b", "b"],
            "group_idx": ("time", [0, 0, 1, 1, 2, 2]),
            "group_idx2": ("time", [0, 1, 1, 1, 1, 1]),
        },
    )

    with xr.set_options(use_flox=use_flox):
        if use_dask:
            ds = ds.chunk()
            if use_lazy_group_idx and module_available("flox", minversion="0.10.5"):
                # This path requires flox installed.
                gs = {
                    g: xr.groupers.UniqueGrouper(labels=np.unique(ds[g]))
                    for g in grp_idx
                }
                actual = getattr(ds.groupby(gs), method)(dim)
            else: