Function test_apply_ufunc_dataarray

xarray/tests/test_units.py:395–415

(variant, dtype)

    ),
)
def test_apply_ufunc_dataarray(variant, dtype):
    variants = {
        "data": (unit_registry.m, 1, 1),
        "dims": (1, unit_registry.m, 1),
        "coords": (1, 1, unit_registry.m),
    }
    data_unit, dim_unit, coord_unit = variants[variant]
    func = functools.partial(
        xr.apply_ufunc, np.mean, input_core_dims=[["x"]], kwargs={"axis": -1}
    )

    array = np.linspace(0, 10, 20).astype(dtype) * data_unit
    x = np.arange(20) * dim_unit
    u = np.linspace(-1, 1, 20) * coord_unit
    data_array = xr.DataArray(data=array, dims="x", coords={"x": x, "u": ("x", u)})

    expected = attach_units(func(strip_units(data_array)), extract_units(data_array))
    actual = func(data_array)

    assert_units_equal(expected, actual)
    assert_identical(expected, actual)


@pytest.mark.parametrize(