From 73fcba607e365b1795c6b1fc937e0d553ab778f0 Mon Sep 17 00:00:00 2001
From: Stephan Hoyer <shoyer@google.com>
Date: Tue, 15 Oct 2024 13:41:48 -0400
Subject: [PATCH] Add missing `memo` argument to DataTree.__deepcopy__

as suggested by @headtr1ck in https://github.com/pydata/xarray/pull/9628#discussion_r1801635289
---
 xarray/core/datatree.py |  8 +++++---
 xarray/core/treenode.py | 23 +++++++++++++++--------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py
index 452a1af2d3b..267911a12a0 100644
--- a/xarray/core/datatree.py
+++ b/xarray/core/datatree.py
@@ -826,12 +826,14 @@ def _replace_node(
 
         self.children = children
 
-    def _copy_node(self, inherit: bool, deep: bool = False) -> Self:
+    def _copy_node(
+        self, inherit: bool, deep: bool = False, memo: dict[int, Any] | None = None
+    ) -> Self:
         """Copy just one node of a tree."""
-        new_node = super()._copy_node(inherit=inherit, deep=deep)
+        new_node = super()._copy_node(inherit=inherit, deep=deep, memo=memo)
         data = self._to_dataset_view(rebuild_dims=False, inherit=inherit)
         if deep:
-            data = data.copy(deep=True)
+            data = data._copy(deep=True, memo=memo)
         new_node._set_node_data(data)
         return new_node
 
diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py
index 102b92dc03b..b28b66cf48e 100644
--- a/xarray/core/treenode.py
+++ b/xarray/core/treenode.py
@@ -274,15 +274,21 @@ def copy(self, *, inherit: bool = True, deep: bool = False) -> Self:
         """
         return self._copy_subtree(inherit=inherit, deep=deep)
 
-    def _copy_subtree(self, inherit: bool, deep: bool = False) -> Self:
+    def _copy_subtree(
+        self, inherit: bool, deep: bool = False, memo: dict[int, Any] | None = None
+    ) -> Self:
         """Copy entire subtree recursively."""
-        new_tree = self._copy_node(inherit=inherit, deep=deep)
+        new_tree = self._copy_node(inherit=inherit, deep=deep, memo=memo)
         for name, child in self.children.items():
             # TODO use `.children[name] = ...` once #9477 is implemented
-            new_tree._set(name, child._copy_subtree(inherit=False, deep=deep))
+            new_tree._set(
+                name, child._copy_subtree(inherit=False, deep=deep, memo=memo)
+            )
         return new_tree
 
-    def _copy_node(self, inherit: bool, deep: bool = False) -> Self:
+    def _copy_node(
+        self, inherit: bool, deep: bool = False, memo: dict[int, Any] | None = None
+    ) -> Self:
         """Copy just one node of a tree"""
         new_empty_node = type(self)()
         return new_empty_node
@@ -291,8 +297,7 @@ def __copy__(self) -> Self:
         return self._copy_subtree(inherit=True, deep=False)
 
     def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Self:
-        del memo  # nodes cannot be reused in a DataTree
-        return self._copy_subtree(inherit=True, deep=True)
+        return self._copy_subtree(inherit=True, deep=True, memo=memo)
 
     def _iter_parents(self: Tree) -> Iterator[Tree]:
         """Iterate up the tree, starting from the current node's parent."""
@@ -693,9 +698,11 @@ def _post_attach(self, parent: Self, name: str) -> None:
         _validate_name(name)  # is this check redundant?
         self._name = name
 
-    def _copy_node(self, inherit: bool, deep: bool = False) -> Self:
+    def _copy_node(
+        self, inherit: bool, deep: bool = False, memo: dict[int, Any] | None = None
+    ) -> Self:
         """Copy just one node of a tree"""
-        new_node = super()._copy_node(inherit=inherit, deep=deep)
+        new_node = super()._copy_node(inherit=inherit, deep=deep, memo=memo)
         new_node._name = self.name
         return new_node