Function convert_units

xarray/tests/test_units.py:230–271

(obj, to)

def convert_units(obj, to):
    # preprocess
    to = {
        key: None if not isinstance(value, unit_registry.Unit) else value
        for key, value in to.items()
    }
    if isinstance(obj, xr.Dataset):
        data_vars = {
            name: convert_units(array.variable, {None: to.get(name)})
            for name, array in obj.data_vars.items()
        }
        coords = {
            name: convert_units(array.variable, {None: to.get(name)})
            for name, array in obj.coords.items()
        }

        new_obj = xr.Dataset(data_vars=data_vars, coords=coords, attrs=obj.attrs)
    elif isinstance(obj, xr.DataArray):
        name = obj.name

        new_units = to.get(name) or to.get("data") or to.get(None) or None
        data = convert_units(obj.variable, {None: new_units})

        coords = {
            name: (array.dims, convert_units(array.variable, {None: to.get(name)}))
            for name, array in obj.coords.items()
            if name != obj.name
        }

        new_obj = xr.DataArray(  # type: ignore[assignment]
            name=name, data=data, coords=coords, attrs=obj.attrs, dims=obj.dims
        )
    elif isinstance(obj, xr.Variable):
        new_data = convert_units(obj.data, to)
        new_obj = obj.copy(data=new_data)  # type: ignore[assignment]
    elif isinstance(obj, unit_registry.Quantity):
        units = to.get(None)
        new_obj = obj.to(units) if units is not None else obj
    else:
        new_obj = obj

    return new_obj


def assert_units_equal(a, b):