|
61 | 61 | import pandas as pd
|
62 | 62 |
|
63 | 63 | from xarray.core.datatree_io import T_DataTreeNetcdfEngine, T_DataTreeNetcdfTypes
|
64 |
| - from xarray.core.merge import CoercibleValue |
| 64 | + from xarray.core.merge import CoercibleMapping, CoercibleValue |
65 | 65 | from xarray.core.types import ErrorOptions, NetcdfWriteModes, ZarrWriteModes
|
66 | 66 |
|
67 | 67 | # """
|
@@ -954,23 +954,29 @@ def update(
|
954 | 954 |
|
955 | 955 | Just like `dict.update` this is an in-place operation.
|
956 | 956 | """
|
957 |
| - # TODO separate by type |
958 | 957 | new_children: dict[str, DataTree] = {}
|
959 |
| - new_variables = {} |
960 |
| - for k, v in other.items(): |
961 |
| - if isinstance(v, DataTree): |
962 |
| - # avoid named node being stored under inconsistent key |
963 |
| - new_child: DataTree = v.copy() |
964 |
| - # Datatree's name is always a string until we fix that (#8836) |
965 |
| - new_child.name = str(k) |
966 |
| - new_children[str(k)] = new_child |
967 |
| - elif isinstance(v, (DataArray, Variable)): |
968 |
| - # TODO this should also accommodate other types that can be coerced into Variables |
969 |
| - new_variables[k] = v |
970 |
| - else: |
971 |
| - raise TypeError(f"Type {type(v)} cannot be assigned to a DataTree") |
972 |
| - |
973 |
| - vars_merge_result = dataset_update_method(self.to_dataset(), new_variables) |
| 958 | + new_variables: CoercibleMapping |
| 959 | + |
| 960 | + if isinstance(other, Dataset): |
| 961 | + new_variables = other |
| 962 | + else: |
| 963 | + new_variables = {} |
| 964 | + for k, v in other.items(): |
| 965 | + if isinstance(v, DataTree): |
| 966 | + # avoid named node being stored under inconsistent key |
| 967 | + new_child: DataTree = v.copy() |
| 968 | + # Datatree's name is always a string until we fix that (#8836) |
| 969 | + new_child.name = str(k) |
| 970 | + new_children[str(k)] = new_child |
| 971 | + elif isinstance(v, (DataArray, Variable)): |
| 972 | + # TODO this should also accommodate other types that can be coerced into Variables |
| 973 | + new_variables[k] = v |
| 974 | + else: |
| 975 | + raise TypeError(f"Type {type(v)} cannot be assigned to a DataTree") |
| 976 | + |
| 977 | + vars_merge_result = dataset_update_method( |
| 978 | + self.to_dataset(inherited=False), new_variables |
| 979 | + ) |
974 | 980 | data = Dataset._construct_direct(**vars_merge_result._asdict())
|
975 | 981 |
|
976 | 982 | # TODO are there any subtleties with preserving order of children like this?
|
|
0 commit comments