Function attach_units

xarray/tests/test_units.py:186–227

(obj, units)

def attach_units(obj, units):
    if not isinstance(obj, xr.DataArray | xr.Dataset | xr.Variable):
        units = units.get("data", None) or units.get(None, None) or 1
        return array_attach_units(obj, units)

    if isinstance(obj, xr.Dataset):
        data_vars = {
            name: attach_units(value, units) for name, value in obj.data_vars.items()
        }

        coords = {
            name: attach_units(value, units) for name, value in obj.coords.items()
        }

        new_obj = xr.Dataset(data_vars=data_vars, coords=coords, attrs=obj.attrs)
    elif isinstance(obj, xr.DataArray):
        # try the array name, "data" and None, then fall back to dimensionless
        data_units = units.get(obj.name, None) or units.get(None, None) or 1

        data = array_attach_units(obj.data, data_units)

        coords = {
            name: (
                (value.dims, array_attach_units(value.data, units.get(name) or 1))
                if name in units
                else (value.dims, value.data)
            )
            for name, value in obj.coords.items()
        }
        dims = obj.dims
        attrs = obj.attrs

        new_obj = xr.DataArray(  # type: ignore[assignment]
            name=obj.name, data=data, coords=coords, attrs=attrs, dims=dims
        )
    else:
        data_units = units.get("data", None) or units.get(None, None) or 1

        data = array_attach_units(obj.data, data_units)
        new_obj = obj.copy(data=data)  # type: ignore[assignment]

    return new_obj


def convert_units(obj, to):