Fix mapping parentage bug (#54)

* think I've fixed the bug

* used feature from python 3.9

* test but doesn't yet work properly

* only check subtree, not down to root

* make sure choice whether to check from root is propagated

* bump python version in CI

* 3.10 instead of 3.1

Author: TomNicholas

.github/workflows/main.yaml

@@ -21,7 +21,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.7, 3.8, 3.9]
+        python-version: ["3.9", "3.10"]
     steps:
       - uses: actions/[email protected]
       - uses: conda-incubator/setup-miniconda@v2
@@ -59,7 +59,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.8, 3.9]
+        python-version: ["3.9", "3.10"]
     steps:
       - uses: actions/[email protected]
       - uses: conda-incubator/setup-miniconda@v2

datatree/_version.py

@@ -1 +1 @@
-__version__ = "0.1.dev75+g977ffe2.d20210902"
+__version__ = "0.1.dev94+g6c6f23c.d20211217"

datatree/mapping.py

@@ -215,6 +215,8 @@ def _map_over_subtree(*args, **kwargs):
         # Find out how many return values we received
         num_return_values = _check_all_return_values(out_data_objects)
 
+        ancestors_of_new_root = first_tree.pathstr.removesuffix(first_tree.name)
+
         # Reconstruct 1+ subtrees from the dict of results, by filling in all nodes of all result trees
         result_trees = []
         for i in range(num_return_values):
@@ -228,7 +230,11 @@ def _map_over_subtree(*args, **kwargs):
                         output_node_data = out_data_objects[p]
                 else:
                     output_node_data = None
-                out_tree_contents[p] = output_node_data
+
+                # Discard parentage so that new trees don't include parents of input nodes
+                # TODO use a proper relative_path method on DataTree(/TreeNode) to do this
+                relative_path = p.removeprefix(ancestors_of_new_root)
+                out_tree_contents[relative_path] = output_node_data
 
             new_tree = DataTree.from_dict(
                 name=first_tree.name, data_objects=out_tree_contents

datatree/testing.py

@@ -41,7 +41,7 @@ def assert_isomorphic(a: DataTree, b: DataTree, from_root: bool = False):
             a = a.root
             b = b.root
 
-        assert a.isomorphic(b, from_root=False), diff_tree_repr(a, b, "isomorphic")
+        assert a.isomorphic(b, from_root=from_root), diff_tree_repr(a, b, "isomorphic")
     else:
         raise TypeError(f"{type(a)} not of type DataTree")
 
@@ -78,7 +78,7 @@ def assert_equal(a: DataTree, b: DataTree, from_root: bool = True):
             a = a.root
             b = b.root
 
-        assert a.equals(b), diff_tree_repr(a, b, "equals")
+        assert a.equals(b, from_root=from_root), diff_tree_repr(a, b, "equals")
     else:
         raise TypeError(f"{type(a)} not of type DataTree")
 
@@ -115,6 +115,6 @@ def assert_identical(a: DataTree, b: DataTree, from_root: bool = True):
             a = a.root
             b = b.root
 
-        assert a.identical(b), diff_tree_repr(a, b, "identical")
+        assert a.identical(b, from_root=from_root), diff_tree_repr(a, b, "identical")
     else:
         raise TypeError(f"{type(a)} not of type DataTree")

datatree/tests/test_mapping.py

@@ -251,6 +251,19 @@ def multiply_then_add(ds, times, add=0.0):
         result_tree = dt.map_over_subtree(multiply_then_add, 10.0, add=2.0)
         assert_equal(result_tree, expected)
 
+    def test_discard_ancestry(self):
+        # Check for datatree GH issue #48
+        dt = create_test_datatree()
+        subtree = dt["set1"]
+
+        @map_over_subtree
+        def times_ten(ds):
+            return 10.0 * ds
+
+        expected = create_test_datatree(modify=lambda ds: 10.0 * ds)["set1"]
+        result_tree = times_ten(subtree)
+        assert_equal(result_tree, expected, from_root=False)
+
 
 @pytest.mark.xfail
 class TestMapOverSubTreeInplace:

setup.py

@@ -26,14 +26,12 @@
         "Topic :: Scientific/Engineering",
         "License :: OSI Approved :: Apache License",
         "Operating System :: OS Independent",
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.7",
-        "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
     ],
     packages=find_packages(exclude=["docs", "tests", "tests.*", "docs.*"]),
     install_requires=install_requires,
-    python_requires=">=3.7",
+    python_requires=">=3.9",
     setup_requires="setuptools_scm",
     use_scm_version={
         "write_to": "datatree/_version.py",