add open_datatree to xarray

doc/roadmap.rst

@@ -156,7 +156,7 @@ types would also be highly useful for xarray users.
 By pursuing these improvements in NumPy we hope to extend the benefits
 to the full scientific Python community, and avoid tight coupling
 between xarray and specific third-party libraries (e.g., for
-implementing untis). This will allow xarray to maintain its domain
+implementing units). This will allow xarray to maintain its domain
 agnostic strengths.
 
 We expect that we may eventually add some minimal interfaces in xarray

pyproject.toml

@@ -96,6 +96,11 @@ warn_redundant_casts = true
 warn_unused_configs = true
 warn_unused_ignores = true
 
+# Ignore mypy errors for modules imported from datatree_.
+[[tool.mypy.overrides]]
+module = "xarray.datatree_.*"
+ignore_errors = true
+
 # Much of the numerical computing stack doesn't have type annotations yet.
 [[tool.mypy.overrides]]
 ignore_missing_imports = true

xarray/backends/api.py

@@ -69,6 +69,7 @@
     T_NetcdfTypes = Literal[
         "NETCDF4", "NETCDF4_CLASSIC", "NETCDF3_64BIT", "NETCDF3_CLASSIC"
     ]
+    from xarray.datatree_.datatree import DataTree
 
 DATAARRAY_NAME = "__xarray_dataarray_name__"
 DATAARRAY_VARIABLE = "__xarray_dataarray_variable__"
@@ -788,6 +789,33 @@ def open_dataarray(
     return data_array
 
 
+def open_datatree(
+    filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
+    engine: T_Engine = None,
+    **kwargs,
+) -> DataTree:
+    """
+    Open and decode a dataset from a file or file-like object, creating one Tree node for each group in the file.
+    Parameters
+    ----------
+    filename_or_obj : str, Path, file-like, or DataStore
+        Strings and Path objects are interpreted as a path to a netCDF file or Zarr store.
+    engine : str, optional
+        Xarray backend engine to us. Valid options include `{"netcdf4", "h5netcdf", "zarr"}`.
+    kwargs :
+        Additional keyword arguments passed to :py:meth:`~xarray.open_dataset` for each group.
+    Returns
+    -------
+    xarray.core.datatree.DataTree
+    """
+    if engine is None:
+        engine = plugins.guess_engine(filename_or_obj)
+
+    backend = plugins.get_backend(engine)
+
+    return backend.open_datatree(filename_or_obj, **kwargs)
+
+
 def open_mfdataset(
     paths: str | NestedSequence[str | os.PathLike],
     chunks: T_Chunks | None = None,

xarray/backends/common.py

@@ -21,6 +21,7 @@
 
     from xarray.core.dataset import Dataset
     from xarray.core.types import NestedSequence
+    from xarray.datatree_.datatree import DataTree
 
 # Create a logger object, but don't add any handlers. Leave that to user code.
 logger = logging.getLogger(__name__)
@@ -458,6 +459,11 @@ class BackendEntrypoint:
     - ``guess_can_open`` method: it shall return ``True`` if the backend is able to open
       ``filename_or_obj``, ``False`` otherwise. The implementation of this
       method is not mandatory.
+    - ``open_datatree`` method: it shall implement reading from file, variables
+      decoding and it returns an instance of :py:class:`~datatree.DataTree`.
+      It shall take in input at least ``filename_or_obj`` argument. The
+      implementation of this method is not mandatory.  For more details see
+      <reference to open_datatree documentation>.
 
     Attributes
     ----------
@@ -508,6 +514,17 @@ def guess_can_open(
 
         return False
 
+    def open_datatree(
+        self,
+        filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
+        **kwargs: Any,
+    ) -> DataTree:
+        """
+        Backend open_datatree method used by Xarray in :py:func:`~xarray.open_datatree`.
+        """
+
+        raise NotImplementedError
+
 
 # mapping of engine name to (module name, BackendEntrypoint Class)
 BACKEND_ENTRYPOINTS: dict[str, tuple[str | None, type[BackendEntrypoint]]] = {}

xarray/backends/h5netcdf_.py

@@ -38,6 +38,7 @@
 
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
+    from xarray.datatree_.datatree import DataTree
 
 
 class H5NetCDFArrayWrapper(BaseNetCDF4Array):
@@ -423,5 +424,45 @@ def open_dataset(  # type: ignore[override]  # allow LSP violation, not supporti
         )
         return ds
 
+    # TODO [MHS, 01/23/2024] This is duplicative of the netcdf4 code in an ugly way.
+    def open_datatree(
+        self,
+        filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
+        **kwargs,
+    ) -> DataTree:
+        from h5netcdf.legacyapi import Dataset as ncDataset
+
+        from xarray.backends.api import open_dataset
+        from xarray.datatree_.datatree import DataTree
+        from xarray.datatree_.datatree.treenode import NodePath
+
+        ds = open_dataset(filename_or_obj, **kwargs)
+        tree_root = DataTree.from_dict({"/": ds})
+        with ncDataset(filename_or_obj, mode="r") as ncds:
+            for path in _iter_nc_groups(ncds):
+                subgroup_ds = open_dataset(filename_or_obj, group=path, **kwargs)
+
+                # TODO refactor to use __setitem__ once creation of new nodes by assigning Dataset works again
+                node_name = NodePath(path).name
+                new_node: DataTree = DataTree(name=node_name, data=subgroup_ds)
+                tree_root._set_item(
+                    path,
+                    new_node,
+                    allow_overwrite=False,
+                    new_nodes_along_path=True,
+                )
+        return tree_root
+
+
+# TODO [MHS, 01/23/2024] directly duplicated from netCDF4 backend
+def _iter_nc_groups(root, parent="/"):
+    from xarray.datatree_.datatree.treenode import NodePath
+
+    parent = NodePath(parent)
+    for path, group in root.groups.items():
+        gpath = parent / path
+        yield str(gpath)
+        yield from _iter_nc_groups(group, parent=gpath)
+
 
 BACKEND_ENTRYPOINTS["h5netcdf"] = ("h5netcdf", H5netcdfBackendEntrypoint)

xarray/backends/netCDF4_.py

@@ -44,6 +44,7 @@
 
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
+    from xarray.datatree_.datatree import DataTree
 
 # This lookup table maps from dtype.byteorder to a readable endian
 # string used by netCDF4.
@@ -667,5 +668,43 @@ def open_dataset(  # type: ignore[override]  # allow LSP violation, not supporti
             )
         return ds
 
+    def open_datatree(
+        self,
+        filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
+        **kwargs,
+    ) -> DataTree:
+        from netCDF4 import Dataset as ncDataset
+
+        from xarray.backends.api import open_dataset
+        from xarray.datatree_.datatree import DataTree
+        from xarray.datatree_.datatree.treenode import NodePath
+
+        ds = open_dataset(filename_or_obj, **kwargs)
+        tree_root = DataTree.from_dict({"/": ds})
+        with ncDataset(filename_or_obj, mode="r") as ncds:
+            for path in _iter_nc_groups(ncds):
+                subgroup_ds = open_dataset(filename_or_obj, group=path, **kwargs)
+
+                # TODO refactor to use __setitem__ once creation of new nodes by assigning Dataset works again
+                node_name = NodePath(path).name
+                new_node: DataTree = DataTree(name=node_name, data=subgroup_ds)
+                tree_root._set_item(
+                    path,
+                    new_node,
+                    allow_overwrite=False,
+                    new_nodes_along_path=True,
+                )
+        return tree_root
+
+
+def _iter_nc_groups(root, parent="/"):
+    from xarray.datatree_.datatree.treenode import NodePath
+
+    parent = NodePath(parent)
+    for path, group in root.groups.items():
+        gpath = parent / path
+        yield str(gpath)
+        yield from _iter_nc_groups(group, parent=gpath)
+
 
 BACKEND_ENTRYPOINTS["netcdf4"] = ("netCDF4", NetCDF4BackendEntrypoint)

xarray/backends/zarr.py

@@ -34,6 +34,7 @@
 
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
+    from xarray.datatree_.datatree import DataTree
 
 
 # need some special secret attributes to tell us the dimensions
@@ -1035,5 +1036,48 @@ def open_dataset(  # type: ignore[override]  # allow LSP violation, not supporti
             )
         return ds
 
+    def open_datatree(
+        self,
+        filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
+        **kwargs,
+    ) -> DataTree:
+        import zarr
+
+        from xarray.backends.api import open_dataset
+        from xarray.datatree_.datatree import DataTree
+        from xarray.datatree_.datatree.treenode import NodePath
+
+        zds = zarr.open_group(filename_or_obj, mode="r")
+        ds = open_dataset(filename_or_obj, engine="zarr", **kwargs)
+        tree_root = DataTree.from_dict({"/": ds})
+        for path in _iter_zarr_groups(zds):
+            try:
+                subgroup_ds = open_dataset(
+                    filename_or_obj, engine="zarr", group=path, **kwargs
+                )
+            except zarr.errors.PathNotFoundError:
+                subgroup_ds = Dataset()
+
+            # TODO refactor to use __setitem__ once creation of new nodes by assigning Dataset works again
+            node_name = NodePath(path).name
+            new_node: DataTree = DataTree(name=node_name, data=subgroup_ds)
+            tree_root._set_item(
+                path,
+                new_node,
+                allow_overwrite=False,
+                new_nodes_along_path=True,
+            )
+        return tree_root
+
+
+def _iter_zarr_groups(root, parent="/"):
+    from xarray.datatree_.datatree.treenode import NodePath
+
+    parent = NodePath(parent)
+    for path, group in root.groups():
+        gpath = parent / path
+        yield str(gpath)
+        yield from _iter_zarr_groups(group, parent=gpath)
+
 
 BACKEND_ENTRYPOINTS["zarr"] = ("zarr", ZarrBackendEntrypoint)

xarray/core/options.py

@@ -20,6 +20,7 @@
         "display_expand_coords",
         "display_expand_data_vars",
         "display_expand_data",
+        "display_expand_groups",
         "display_expand_indexes",
         "display_default_indexes",
         "enable_cftimeindex",
@@ -44,6 +45,7 @@ class T_Options(TypedDict):
         display_expand_coords: Literal["default", True, False]
         display_expand_data_vars: Literal["default", True, False]
         display_expand_data: Literal["default", True, False]
+        display_expand_groups: Literal["default", True, False]
         display_expand_indexes: Literal["default", True, False]
         display_default_indexes: Literal["default", True, False]
         enable_cftimeindex: bool
@@ -68,6 +70,7 @@ class T_Options(TypedDict):
     "display_expand_coords": "default",
     "display_expand_data_vars": "default",
     "display_expand_data": "default",
+    "display_expand_groups": "default",
     "display_expand_indexes": "default",
     "display_default_indexes": False,
     "enable_cftimeindex": True,

xarray/datatree_/datatree/__init__.py

@@ -1,25 +1,15 @@
 # import public API
 from .datatree import DataTree
 from .extensions import register_datatree_accessor
-from .io import open_datatree
 from .mapping import TreeIsomorphismError, map_over_subtree
 from .treenode import InvalidTreeError, NotFoundInTreeError
 
-try:
-    # NOTE: the `_version.py` file must not be present in the git repository
-    #   as it is generated by setuptools at install time
-    from ._version import __version__
-except ImportError:  # pragma: no cover
-    # Local copy or not installed with setuptools
-    __version__ = "999"
 
 __all__ = (
     "DataTree",
-    "open_datatree",
     "TreeIsomorphismError",
     "InvalidTreeError",
     "NotFoundInTreeError",
     "map_over_subtree",
     "register_datatree_accessor",
-    "__version__",
 )

xarray/datatree_/datatree/datatree.py

@@ -16,6 +16,7 @@
     List,
     Mapping,
     MutableMapping,
+    NoReturn,
     Optional,
     Set,
     Tuple,
@@ -160,7 +161,7 @@ def __setitem__(self, key, val) -> None:
             "use `.copy()` first to get a mutable version of the input dataset."
         )
 
-    def update(self, other) -> None:
+    def update(self, other) -> NoReturn:
         raise AttributeError(
             "Mutation of the DatasetView is not allowed, please use `.update` on the wrapping DataTree node, "
             "or use `dt.to_dataset()` if you want a mutable dataset. If calling this from within `map_over_subtree`,"

xarray/datatree_/datatree/formatting_html.py

@@ -10,9 +10,6 @@
     datavar_section,
     dim_section,
 )
-from xarray.core.options import OPTIONS
-
-OPTIONS["display_expand_groups"] = "default"
 
 
 def summarize_children(children: Mapping[str, Any]) -> str: