Function merge_trees

xarray/structure/merge.py:803–851

Merge specialized to DataTree objects.

(
    trees: Sequence[DataTree],
    compat: CompatOptions | CombineKwargDefault = _COMPAT_DEFAULT,
    join: JoinOptions | CombineKwargDefault = _JOIN_DEFAULT,
    fill_value: object = dtypes.NA,
    combine_attrs: CombineAttrsOptions = "override",
)

def merge_trees(
    trees: Sequence[DataTree],
    compat: CompatOptions | CombineKwargDefault = _COMPAT_DEFAULT,
    join: JoinOptions | CombineKwargDefault = _JOIN_DEFAULT,
    fill_value: object = dtypes.NA,
    combine_attrs: CombineAttrsOptions = "override",
) -> DataTree:
    """Merge specialized to DataTree objects."""
    from xarray.core.dataset import Dataset
    from xarray.core.datatree import DataTree
    from xarray.core.datatree_mapping import add_path_context_to_errors

    if fill_value is not dtypes.NA:
        # fill_value support dicts, which probably should be mapped to sub-groups?
        raise NotImplementedError(
            "fill_value is not yet supported for DataTree objects in merge"
        )

    node_lists: defaultdict[str, list[DataTree]] = defaultdict(list)
    for tree in trees:
        for key, node in tree.subtree_with_keys:
            node_lists[key].append(node)

    root_datasets = [node.dataset for node in node_lists.pop(".")]
    with add_path_context_to_errors("."):
        root_ds = merge(
            root_datasets, compat=compat, join=join, combine_attrs=combine_attrs
        )
    result = DataTree(dataset=root_ds)

    def level(kv):
        # all trees with the same path have the same level
        _, trees = kv
        return trees[0].level

    for key, nodes in sorted(node_lists.items(), key=level):
        # Merge datasets, including inherited indexes to ensure alignment.
        datasets = [node.dataset for node in nodes]
        with add_path_context_to_errors(key):
            merge_result = merge_core(
                datasets,
                compat=compat,
                join=join,
                combine_attrs=combine_attrs,
            )
        merged_ds = Dataset._construct_direct(**merge_result._asdict())
        result[key] = DataTree(dataset=merged_ds)

    return result


@overload