Fix inadvertent deep-copying of child data in DataTree (pydata#9684)

headtr1ck · Oct 26, 2024 · dbb98b4 · dbb98b4
1 parent 5be821b
commit dbb98b4
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 1 deletion.
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
@@ -34,6 +34,8 @@ Deprecations
 Bug fixes
 ~~~~~~~~~
 
+- Fix inadvertent deep-copying of child data in DataTree.
+  By `Stephan Hoyer <https://github.com/shoyer>`_.
 
 Documentation
 ~~~~~~~~~~~~~

diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py
@@ -146,7 +146,7 @@ def check_alignment(
 ) -> None:
     if parent_ds is not None:
         try:
-            align(node_ds, parent_ds, join="exact")
+            align(node_ds, parent_ds, join="exact", copy=False)
         except ValueError as e:
             node_repr = _indented(_without_header(repr(node_ds)))
             parent_repr = _indented(dims_and_coords_repr(parent_ds))

diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py
@@ -69,6 +69,19 @@ def test_data_arg(self) -> None:
         with pytest.raises(TypeError):
             DataTree(dataset=xr.DataArray(42, name="foo"))  # type: ignore[arg-type]
 
+    def test_child_data_not_copied(self) -> None:
+        # regression test for https://github.com/pydata/xarray/issues/9683
+        class NoDeepCopy:
+            def __deepcopy__(self, memo):
+                raise TypeError("class can't be deepcopied")
+
+        da = xr.DataArray(NoDeepCopy())
+        ds = xr.Dataset({"var": da})
+        dt1 = xr.DataTree(ds)
+        dt2 = xr.DataTree(ds, children={"child": dt1})
+        dt3 = xr.DataTree.from_dict({"/": ds, "child": ds})
+        assert_identical(dt2, dt3)
+
 
 class TestFamilyTree:
     def test_dont_modify_children_inplace(self) -> None: