Function _broadcast_helper

xarray/structure/alignment.py:1136–1172

(
    arg: T_Alignable, exclude, dims_map, common_coords
)

def _broadcast_helper(
    arg: T_Alignable, exclude, dims_map, common_coords
) -> T_Alignable:
    from xarray.core.dataarray import DataArray
    from xarray.core.dataset import Dataset

    def _set_dims(var):
        # Add excluded dims to a copy of dims_map
        var_dims_map = dims_map.copy()
        for dim in exclude:
            with suppress(ValueError):
                # ignore dim not in var.dims
                var_dims_map[dim] = var.shape[var.dims.index(dim)]

        return var.set_dims(var_dims_map)

    def _broadcast_array(array: T_DataArray) -> T_DataArray:
        data = _set_dims(array.variable)
        coords = dict(array.coords)
        coords.update(common_coords)
        return array.__class__(
            data, coords, data.dims, name=array.name, attrs=array.attrs
        )

    def _broadcast_dataset(ds: T_Dataset) -> T_Dataset:
        data_vars = {k: _set_dims(ds.variables[k]) for k in ds.data_vars}
        coords = dict(ds.coords)
        coords.update(common_coords)
        return ds.__class__(data_vars, coords, ds.attrs)

    # remove casts once https://github.com/python/mypy/issues/12800 is resolved
    if isinstance(arg, DataArray):
        return _broadcast_array(arg)  # type: ignore[return-value,unused-ignore]
    elif isinstance(arg, Dataset):
        return _broadcast_dataset(arg)  # type: ignore[return-value,unused-ignore]
    else:
        raise ValueError("all input must be Dataset or DataArray objects")


@overload