Function collect_variables_and_indexes

xarray/structure/merge.py:358–418

Collect variables and indexes from list of mappings of xarray objects. Mappings can be Dataset or Coordinates objects, in which case both variables and indexes are extracted from it. It can also have values of one of the following types: - an xarray.Variable - a tuple `(dims, d

(
    list_of_mappings: Iterable[DatasetLike],
    indexes: Mapping[Any, Any] | None = None,
)

def collect_variables_and_indexes(
    list_of_mappings: Iterable[DatasetLike],
    indexes: Mapping[Any, Any] | None = None,
) -> dict[Hashable, list[MergeElement]]:
    """Collect variables and indexes from list of mappings of xarray objects.

    Mappings can be Dataset or Coordinates objects, in which case both
    variables and indexes are extracted from it.

    It can also have values of one of the following types:
    - an xarray.Variable
    - a tuple `(dims, data[, attrs[, encoding]])` that can be converted in
      an xarray.Variable
    - or an xarray.DataArray

    If a mapping of indexes is given, those indexes are assigned to all variables
    with a matching key/name. For dimension variables with no matching index, a
    default (pandas) index is assigned. DataArray indexes that don't match mapping
    keys are also extracted.

    """
    from xarray.core.coordinates import Coordinates
    from xarray.core.dataarray import DataArray
    from xarray.core.dataset import Dataset

    if indexes is None:
        indexes = {}

    grouped: dict[Hashable, list[MergeElement]] = defaultdict(list)

    def append(name, variable, index):
        grouped[name].append((variable, index))

    def append_all(variables, indexes):
        for name, variable in variables.items():
            append(name, variable, indexes.get(name))

    for mapping in list_of_mappings:
        if isinstance(mapping, Coordinates | Dataset):
            append_all(mapping.variables, mapping.xindexes)
            continue

        for name, variable in mapping.items():
            if isinstance(variable, DataArray):
                coords_ = variable._coords.copy()  # use private API for speed
                indexes_ = dict(variable._indexes)
                # explicitly overwritten variables should take precedence
                coords_.pop(name, None)
                indexes_.pop(name, None)
                append_all(coords_, indexes_)

            variable = as_variable(variable, name=name, auto_convert=False)
            if name in indexes:
                append(name, variable, indexes[name])
            elif variable.dims == (name,):
                idx, idx_vars = create_default_index_implicit(variable)
                append_all(idx_vars, dict.fromkeys(idx_vars, idx))
            else: