Function test_align_dataset

xarray/tests/test_units.py:587–664

(value, unit, variant, error, dtype)

)
@pytest.mark.parametrize("value", (10, dtypes.NA))
def test_align_dataset(value, unit, variant, error, dtype):
    if variant == "coords" and (
        value != dtypes.NA or isinstance(unit, unit_registry.Unit)
    ):
        pytest.xfail(
            reason=(
                "fill_value is used for both data variables and coords. "
                "See https://github.com/pydata/xarray/issues/4165"
            )
        )

    fill_value = dtypes.get_fill_value(dtype) if value == dtypes.NA else value

    original_unit = unit_registry.m

    variants = {
        "data": ((original_unit, unit), (1, 1), (1, 1)),
        "dims": ((1, 1), (original_unit, unit), (1, 1)),
        "coords": ((1, 1), (1, 1), (original_unit, unit)),
    }
    (
        (data_unit1, data_unit2),
        (dim_unit1, dim_unit2),
        (coord_unit1, coord_unit2),
    ) = variants[variant]

    array1 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * data_unit1
    array2 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * data_unit2

    x = np.arange(2) * dim_unit1
    y1 = np.arange(5) * dim_unit1
    y2 = np.arange(2, 7) * dim_unit2

    u1 = np.array([3, 5, 7, 8, 9]) * coord_unit1
    u2 = np.array([7, 8, 9, 11, 13]) * coord_unit2

    coords1 = {"x": x, "y": y1}
    coords2 = {"x": x, "y": y2}
    if variant == "coords":
        coords1["u"] = ("y", u1)
        coords2["u"] = ("y", u2)

    ds1 = xr.Dataset(data_vars={"a": (("x", "y"), array1)}, coords=coords1)
    ds2 = xr.Dataset(data_vars={"a": (("x", "y"), array2)}, coords=coords2)

    fill_value = fill_value * data_unit2
    func = function(xr.align, join="outer", fill_value=fill_value)
    if error is not None and (value != dtypes.NA or isinstance(fill_value, Quantity)):
        with pytest.raises(error):
            func(ds1, ds2)

        return

    stripped_kwargs = {
        key: strip_units(
            convert_units(value, {None: data_unit1 if data_unit2 != 1 else None})
        )
        for key, value in func.kwargs.items()