Refine DataTree.equals and DataTree.identical

xarray/core/datatree.py

@@ -1238,65 +1238,67 @@ def isomorphic(
         except (TypeError, TreeIsomorphismError):
             return False
 
-    def equals(self, other: DataTree, from_root: bool = True) -> bool:
+    def _matching(
+        self,
+        other: DataTree,
+        nodes_match: Callable[[DataTree, DataTree], bool],
+    ) -> bool:
+        if not self.isomorphic(other, from_root=True, strict_names=True):
+            return False
+
+        return all(
+            nodes_match(node, other_node)
+            for node, other_node in zip(self.descendants, other.descendants)
+        )
+
+    def equals(self, other: DataTree) -> bool:
         """
         Two DataTrees are equal if they have isomorphic node structures, with matching node names,
         and if they have matching variables and coordinates, all of which are equal.
 
-        By default this method will check the whole tree above the given node.
-
         Parameters
         ----------
         other : DataTree
             The other tree object to compare to.
-        from_root : bool, optional, default is True
-            Whether or not to first traverse to the root of the two trees before checking for isomorphism.
-            If neither tree has a parent then this has no effect.
 
         See Also
         --------
         Dataset.equals
         DataTree.isomorphic
         DataTree.identical
         """
-        if not self.isomorphic(other, from_root=from_root, strict_names=True):
+        to_ds = lambda x: x._to_dataset_view(inherited=True, rebuild_dims=False)
+        matcher = lambda x, y: to_ds(x).equals(to_ds(y))
+        if not matcher(self, other):
             return False
+        return self._matching(other, nodes_match=matcher)
 
-        return all(
-            [
-                node.ds.equals(other_node.ds)
-                for node, other_node in zip(self.subtree, other.subtree)
-            ]
-        )
-
-    def identical(self, other: DataTree, from_root=True) -> bool:
+    def identical(self, other: DataTree) -> bool:
         """
         Like equals, but will also check all dataset attributes and the attributes on
-        all variables and coordinates.
-
-        By default this method will check the whole tree above the given node.
+        all variables and coordinates, and requires the coordinates are defined at the
+        exact same levels of the DataTree hierarchy.
 
         Parameters
         ----------
         other : DataTree
             The other tree object to compare to.
-        from_root : bool, optional, default is True
-            Whether or not to first traverse to the root of the two trees before checking for isomorphism.
-            If neither tree has a parent then this has no effect.
 
         See Also
         --------
         Dataset.identical
         DataTree.isomorphic
         DataTree.equals
         """
-        if not self.isomorphic(other, from_root=from_root, strict_names=True):
+        # include inherited coordinates only at the root; otherwise require that
+        # everything is also defined at the same level
+        self_view = self._to_dataset_view(inherited=True, rebuild_dims=False)
+        other_view = other._to_dataset_view(inherited=True, rebuild_dims=False)
+        if not self_view.identical(other_view):
             return False
-
-        return all(
-            node.ds.identical(other_node.ds)
-            for node, other_node in zip(self.subtree, other.subtree)
-        )
+        to_ds = lambda x: x._to_dataset_view(inherited=False, rebuild_dims=False)
+        matcher = lambda x, y: to_ds(x).identical(to_ds(y))
+        return self._matching(other, nodes_match=matcher)
 
     def filter(self: DataTree, filterfunc: Callable[[DataTree], bool]) -> DataTree:
         """

xarray/testing/assertions.py

@@ -116,7 +116,7 @@ def assert_equal(a: DataTree, b: DataTree, from_root: bool = True): ...
 
 
 @ensure_warnings
-def assert_equal(a, b, from_root=True, check_dim_order: bool = True):
+def assert_equal(a, b, check_dim_order: bool = True):
     """Like :py:func:`numpy.testing.assert_array_equal`, but for xarray
     objects.
 
@@ -135,10 +135,6 @@ def assert_equal(a, b, from_root=True, check_dim_order: bool = True):
         or xarray.core.datatree.DataTree. The first object to compare.
     b : xarray.Dataset, xarray.DataArray, xarray.Variable, xarray.Coordinates
         or xarray.core.datatree.DataTree. The second object to compare.
-    from_root : bool, optional, default is True
-        Only used when comparing DataTree objects. Indicates whether or not to
-        first traverse to the root of the trees before checking for isomorphism.
-        If a & b have no parents then this has no effect.
     check_dim_order : bool, optional, default is True
         Whether dimensions must be in the same order.
 
@@ -159,11 +155,7 @@ def assert_equal(a, b, from_root=True, check_dim_order: bool = True):
     elif isinstance(a, Coordinates):
         assert a.equals(b), formatting.diff_coords_repr(a, b, "equals")
     elif isinstance(a, DataTree):
-        if from_root:
-            a = a.root
-            b = b.root
-
-        assert a.equals(b, from_root=from_root), diff_datatree_repr(a, b, "equals")
+        assert a.equals(b), diff_datatree_repr(a, b, "equals")
     else:
         raise TypeError(f"{type(a)} not supported by assertion comparison")
 
@@ -173,11 +165,11 @@ def assert_identical(a, b): ...
 
 
 @overload
-def assert_identical(a: DataTree, b: DataTree, from_root: bool = True): ...
+def assert_identical(a: DataTree, b: DataTree): ...
 
 
 @ensure_warnings
-def assert_identical(a, b, from_root=True):
+def assert_identical(a, b):
     """Like :py:func:`xarray.testing.assert_equal`, but also matches the
     objects' names and attributes.
 
@@ -193,10 +185,6 @@ def assert_identical(a, b, from_root=True):
         The first object to compare.
     b : xarray.Dataset, xarray.DataArray, xarray.Variable or xarray.Coordinates
         The second object to compare.
-    from_root : bool, optional, default is True
-        Only used when comparing DataTree objects. Indicates whether or not to
-        first traverse to the root of the trees before checking for isomorphism.
-        If a & b have no parents then this has no effect.
     check_dim_order : bool, optional, default is True
         Whether dimensions must be in the same order.
 
@@ -220,13 +208,7 @@ def assert_identical(a, b, from_root=True):
     elif isinstance(a, Coordinates):
         assert a.identical(b), formatting.diff_coords_repr(a, b, "identical")
     elif isinstance(a, DataTree):
-        if from_root:
-            a = a.root
-            b = b.root
-
-        assert a.identical(b, from_root=from_root), diff_datatree_repr(
-            a, b, "identical"
-        )
+        assert a.identical(b), diff_datatree_repr(a, b, "identical")
     else:
         raise TypeError(f"{type(a)} not supported by assertion comparison")
 

xarray/tests/test_datatree.py

@@ -1091,6 +1091,62 @@ def f(x, tree, y):
         assert actual is dt and actual.attrs == attrs
 
 
+class TestEqualsAndIdentical:
+    def test_basics(self, create_test_datatree):
+        dt = create_test_datatree()
+        assert dt.equals(dt)
+        assert dt.identical(dt)
+
+        other = DataTree()
+        assert not dt.equals(other)
+        assert not dt.identical(other)
+        assert not other.equals(dt)
+        assert not other.identical(dt)
+
+    def test_isomormorphic(self, create_test_datatree):
+        dt = create_test_datatree()
+
+        dt_new_node = dt.copy()
+        dt_new_node["/something_new"] = DataTree()
+        assert not dt.equals(dt_new_node)
+        assert not dt.identical(dt_new_node)
+        assert not dt_new_node.equals(dt)
+        assert not dt_new_node.identical(dt)
+
+        dt_new_array = dt.copy()
+        dt_new_array["/something_else"] = xr.DataArray(1234)
+        assert not dt.equals(dt_new_array)
+        assert not dt.identical(dt_new_array)
+        assert not dt_new_array.equals(dt)
+        assert not dt_new_array.identical(dt)
+
+    def test_equal_but_not_identical_inheritance(self):
+        duplicated_coords = DataTree.from_dict(
+            {
+                "/": Dataset(coords={"x": [1]}),
+                "/sub": Dataset(coords={"x": [1]}),
+            }
+        )
+        inherited_coords = DataTree.from_dict(
+            {
+                "/": Dataset(coords={"x": [1]}),
+                "/sub": Dataset(),
+            }
+        )
+        assert duplicated_coords.equals(inherited_coords)
+        assert inherited_coords.equals(duplicated_coords)
+        assert not duplicated_coords.identical(inherited_coords)
+        assert not inherited_coords.identical(duplicated_coords)
+
+    def test_equal_but_not_identical_attrs(self):
+        without_attrs = DataTree()
+        with_attrs = DataTree(Dataset(attrs={"foo": "bar"}))
+        assert without_attrs.equals(with_attrs)
+        assert with_attrs.equals(without_attrs)
+        assert not without_attrs.identical(with_attrs)
+        assert not with_attrs.identical(without_attrs)
+
+
 class TestSubset:
     def test_match(self):
         # TODO is this example going to cause problems with case sensitivity?

xarray/tests/test_datatree_mapping.py

@@ -251,9 +251,9 @@ def test_discard_ancestry(self, create_test_datatree):
         def times_ten(ds):
             return 10.0 * ds
 
-        expected = create_test_datatree(modify=lambda ds: 10.0 * ds)["set1"]
+        expected = create_test_datatree(modify=lambda ds: 10.0 * ds)["set1"].copy()
         result_tree = times_ten(subtree)
-        assert_equal(result_tree, expected, from_root=False)
+        assert_equal(result_tree, expected)
 
     def test_skip_empty_nodes_with_attrs(self, create_test_datatree):
         # inspired by xarray-datatree GH262