Skip to content
This repository was archived by the owner on Oct 24, 2024. It is now read-only.

to/from_dict #82

Merged
merged 5 commits into from
May 5, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions datatree/datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
TYPE_CHECKING,
Any,
Callable,
Dict,
Generic,
Iterable,
Mapping,
Expand Down Expand Up @@ -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:
"""
Expand All @@ -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,
Expand All @@ -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)
Expand Down
14 changes: 14 additions & 0 deletions datatree/tests/test_datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
...
Expand Down
2 changes: 1 addition & 1 deletion docs/source/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions docs/source/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
~~~~~~~~~~~~~~~~
Expand Down