to/from_dict xarray-contrib/datatree#82

* add to_dict method

* added roundtrip test

* add xfailed test for rountripping with named root

* whats-new

Author: TomNicholas

datatree/datatree.py

@@ -6,6 +6,7 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    Dict,
     Generic,
     Iterable,
     Mapping,
@@ -316,7 +317,7 @@ def update(self, other: Dataset | Mapping[str, DataTree | DataArray]) -> None:
     @classmethod
     def from_dict(
         cls,
-        d: MutableMapping[str, DataTree | Dataset | DataArray],
+        d: MutableMapping[str, Dataset | DataArray | None],
         name: str = None,
     ) -> DataTree:
         """
@@ -337,19 +338,22 @@ def from_dict(
         Returns
         -------
         DataTree
+
+        Notes
+        -----
+        If your dictionary is nested you will need to flatten it before using this method.
         """
 
         # First create the root node
-        # TODO there is a real bug here where what if root_data is of type DataTree?
         root_data = d.pop("/", None)
-        obj = cls(name=name, data=root_data, parent=None, children=None)  # type: ignore[arg-type]
+        obj = cls(name=name, data=root_data, parent=None, children=None)
 
         if d:
             # Populate tree with children determined from data_objects mapping
             for path, data in d.items():
                 # Create and set new node
                 node_name = NodePath(path).name
-                new_node = cls(name=node_name, data=data)  # type: ignore[arg-type]
+                new_node = cls(name=node_name, data=data)
                 obj._set_item(
                     path,
                     new_node,
@@ -358,6 +362,16 @@ def from_dict(
                 )
         return obj
 
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Create a dictionary mapping of absolute node paths to the data contained in those nodes.
+
+        Returns
+        -------
+        Dict
+        """
+        return {node.path: node.ds for node in self.subtree}
+
     @property
     def nbytes(self) -> int:
         return sum(node.ds.nbytes if node.has_data else 0 for node in self.subtree)

datatree/tests/test_datatree.py

@@ -328,6 +328,20 @@ def test_full(self):
             "/set3",
         ]
 
+    def test_roundtrip(self):
+        dt = create_test_datatree()
+        roundtrip = DataTree.from_dict(dt.to_dict())
+        assert roundtrip.equals(dt)
+
+    @pytest.mark.xfail
+    def test_roundtrip_unnamed_root(self):
+        # See GH81
+
+        dt = create_test_datatree()
+        dt.name = "root"
+        roundtrip = DataTree.from_dict(dt.to_dict())
+        assert roundtrip.equals(dt)
+
 
 class TestBrowsing:
     ...

docs/source/api.rst

@@ -215,13 +215,13 @@ I/O
 
    open_datatree
    DataTree.from_dict
+   DataTree.to_dict
    DataTree.to_netcdf
    DataTree.to_zarr
 
 ..
 
    Missing
-   DataTree.to_dict
    open_mfdatatree
 
 Exceptions

docs/source/whats-new.rst

@@ -38,6 +38,8 @@ New Features
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
 - New delitem method so you can delete nodes. (:pull:`88`)
   By `Tom Nicholas <https://github.com/TomNicholas>`_.
+- New ``to_dict`` method. (:pull:`82`)
+  By `Tom Nicholas <https://github.com/TomNicholas>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~