From 5ae7a47c7770f6ce5f2f36fa06ffcea1c60e1c58 Mon Sep 17 00:00:00 2001 From: Simon Perkins Date: Tue, 15 Jul 2025 19:25:19 +0200 Subject: [PATCH] Preserve encoding in DataTree.assign --- xarray/core/datatree.py | 1 + xarray/tests/test_datatree.py | 8 ++++++++ 2 files changed, 9 insertions(+) diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index afef2f20094..ae34884bf6d 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -1060,6 +1060,7 @@ def update( self.to_dataset(inherit=False), new_variables ) data = Dataset._construct_direct(**vars_merge_result._asdict()) + data.encoding = self.encoding # TODO are there any subtleties with preserving order of children like this? merged_children = {**self.children, **new_children} diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 2bf079a7cbd..e0834518a6b 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -1679,6 +1679,14 @@ def test_assign(self) -> None: result = dt.assign({"foo": xr.DataArray(0), "a": DataTree()}) assert_equal(result, expected) + def test_assign_encoding(self) -> None: + ds = xr.Dataset({"foo": 0}) + ds.encoding = {"test": "test"} + node = xr.DataTree(ds) + assert node.encoding == {"test": "test"} + node = node.assign(bar=xr.DataArray(1)) + assert node.encoding == {"test": "test"} + def test_filter_like(self) -> None: flower_tree = DataTree.from_dict( {"root": None, "trunk": None, "leaves": None, "flowers": None}