Function test_grid_rechunk

xarray/tests/test_backends_chunks.py:98–116

(enc_chunks, region, nd_v_chunks, expected_chunks)

    ],
)
def test_grid_rechunk(enc_chunks, region, nd_v_chunks, expected_chunks):
    dims = [f"dim_{i}" for i in range(len(region))]
    coords = {
        dim: list(range(r.start, r.stop)) for dim, r in zip(dims, region, strict=False)
    }
    shape = tuple(r.stop - r.start for r in region)
    arr = xr.DataArray(
        np.arange(np.prod(shape)).reshape(shape),
        dims=dims,
        coords=coords,
    )
    arr = arr.chunk(dict(zip(dims, nd_v_chunks, strict=False)))

    result = grid_rechunk(
        arr.variable,
        enc_chunks=enc_chunks,
        region=region,
    )
    assert result.chunks == expected_chunks