Merge xarray-contrib/datatree#14 from TomNicholas/all_properties

Define all Dataset properties on DataTree

Author: TomNicholas

datatree/datatree.py

@@ -92,11 +92,15 @@ def _map_over_subtree(tree, *args, **kwargs):
     return _map_over_subtree
 
 
+_DATASET_PROPERTIES_TO_EXPOSE = ['dims', 'variables', 'encoding', 'sizes', 'attrs', 'nbytes', 'indexes', 'xindexes',
+                                 'xindexes', 'coords', 'data_vars', 'chunks', 'real', 'imag']
+
+
 class DatasetPropertiesMixin:
     """Expose properties of wrapped Dataset"""
 
-    # TODO a neater / more succinct way of doing this?
-    # we wouldn't need it at all if we inherited directly from Dataset...
+    # TODO a neater way of setting all of these?
+    # We wouldn't need this at all if we inherited directly from Dataset...
 
     @property
     def dims(self):
@@ -133,13 +137,73 @@ def attrs(self):
         else:
             raise AttributeError("property is not defined for a node with no data")
 
+
+    @property
+    def nbytes(self) -> int:
+        return sum(node.ds.nbytes for node in self.subtree_nodes)
+
+    @property
+    def indexes(self):
+        if self.has_data:
+            return self.ds.indexes
+        else:
+            raise AttributeError("property is not defined for a node with no data")
+
+    @property
+    def xindexes(self):
+        if self.has_data:
+            return self.ds.xindexes
+        else:
+            raise AttributeError("property is not defined for a node with no data")
+
+    @property
+    def coords(self):
+        if self.has_data:
+            return self.ds.coords
+        else:
+            raise AttributeError("property is not defined for a node with no data")
+
+    @property
+    def data_vars(self):
+        if self.has_data:
+            return self.ds.data_vars
+        else:
+            raise AttributeError("property is not defined for a node with no data")
+
+    # TODO should this instead somehow give info about the chunking of every node?
+    @property
+    def chunks(self):
+        if self.has_data:
+            return self.ds.chunks
+        else:
+            raise AttributeError("property is not defined for a node with no data")
+
+    @property
+    def real(self):
+        if self.has_data:
+            return self.ds.real
+        else:
+            raise AttributeError("property is not defined for a node with no data")
+
+    @property
+    def imag(self):
+        if self.has_data:
+            return self.ds.imag
+        else:
+            raise AttributeError("property is not defined for a node with no data")
+
     # TODO .loc
 
     dims.__doc__ = Dataset.dims.__doc__
     variables.__doc__ = Dataset.variables.__doc__
     encoding.__doc__ = Dataset.encoding.__doc__
     sizes.__doc__ = Dataset.sizes.__doc__
     attrs.__doc__ = Dataset.attrs.__doc__
+    indexes.__doc__ = Dataset.indexes.__doc__
+    xindexes.__doc__ = Dataset.xindexes.__doc__
+    coords.__doc__ = Dataset.coords.__doc__
+    data_vars.__doc__ = Dataset.data_vars.__doc__
+    chunks.__doc__ = Dataset.chunks.__doc__
 
 
 _MAPPED_DOCSTRING_ADDENDUM = textwrap.fill("This method was copied from xarray.Dataset, but has been altered to "
@@ -283,7 +347,7 @@ def __init__(
         self._add_all_dataset_api()
 
     def _add_all_dataset_api(self):
-        # Add methods like .mean(), but wrapped to map over subtrees
+        # Add methods like .isel(), but wrapped to map over subtrees
         self._add_dataset_methods()
 
         # TODO add dataset ops here
@@ -633,7 +697,7 @@ def merge_child_datasets(
         datasets = [self.get(path).ds for path in paths]
         return merge(datasets, compat=compat, join=join, fill_value=fill_value, combine_attrs=combine_attrs)
 
-    def as_dataarray(self) -> DataArray:
+    def as_array(self) -> DataArray:
         return self.ds.as_dataarray()
 
     @property

datatree/tests/test_dataset_api.py

@@ -80,6 +80,7 @@ def test_properties(self):
         assert dt.sizes == dt.ds.sizes
         assert dt.variables == dt.ds.variables
 
+
     def test_no_data_no_properties(self):
         dt = DataNode('root', data=None)
         with pytest.raises(AttributeError):