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

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
+      :ref:`RST open_datatree`.
 
     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,41 @@ 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: str, **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, **kwargs)
+        tree_root = DataTree.from_dict({"/": ds})
+        with ncDataset(filename, mode="r") as ncds:
+            for path in _iter_nc_groups(ncds):
+                subgroup_ds = open_dataset(filename, 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,39 @@ def open_dataset(  # type: ignore[override]  # allow LSP violation, not supporti
             )
         return ds
 
+    def open_datatree(self, filename: str, **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, **kwargs)
+        tree_root = DataTree.from_dict({"/": ds})
+        with ncDataset(filename, mode="r") as ncds:
+            for path in _iter_nc_groups(ncds):
+                subgroup_ds = open_dataset(filename, 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,42 @@ def open_dataset(  # type: ignore[override]  # allow LSP violation, not supporti
             )
         return ds
 
+    def open_datatree(self, store, **kwargs) -> DataTree:
+        import zarr  # type: ignore
+
+        from xarray.backends.api import open_dataset
+        from xarray.datatree_.datatree import DataTree
+        from xarray.datatree_.datatree.treenode import NodePath
+
+        zds = zarr.open_group(store, mode="r")
+        ds = open_dataset(store, engine="zarr", **kwargs)
+        tree_root = DataTree.from_dict({"/": ds})
+        for path in _iter_zarr_groups(zds):
+            try:
+                subgroup_ds = open_dataset(store, 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/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/io.py

@@ -1,22 +1,4 @@
-from xarray import Dataset, open_dataset
-
-from .datatree import DataTree, NodePath
-
-
-def _iter_zarr_groups(root, parent="/"):
-    parent = NodePath(parent)
-    for path, group in root.groups():
-        gpath = parent / path
-        yield str(gpath)
-        yield from _iter_zarr_groups(group, parent=gpath)
-
-
-def _iter_nc_groups(root, parent="/"):
-    parent = NodePath(parent)
-    for path, group in root.groups.items():
-        gpath = parent / path
-        yield str(gpath)
-        yield from _iter_nc_groups(group, parent=gpath)
+from .datatree import DataTree
 
 
 def _get_nc_dataset_class(engine):
@@ -34,77 +16,6 @@ def _get_nc_dataset_class(engine):
     return Dataset
 
 
-def open_datatree(filename_or_obj, 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 ``xarray.open_dataset`` for each group.
-
-    Returns
-    -------
-    DataTree
-    """
-
-    if engine == "zarr":
-        return _open_datatree_zarr(filename_or_obj, **kwargs)
-    elif engine in [None, "netcdf4", "h5netcdf"]:
-        return _open_datatree_netcdf(filename_or_obj, engine=engine, **kwargs)
-    else:
-        raise ValueError("Unsupported engine")
-
-
-def _open_datatree_netcdf(filename: str, **kwargs) -> DataTree:
-    ncDataset = _get_nc_dataset_class(kwargs.get("engine", None))
-
-    ds = open_dataset(filename, **kwargs)
-    tree_root = DataTree.from_dict({"/": ds})
-    with ncDataset(filename, mode="r") as ncds:
-        for path in _iter_nc_groups(ncds):
-            subgroup_ds = open_dataset(filename, 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 _open_datatree_zarr(store, **kwargs) -> DataTree:
-    import zarr  # type: ignore
-
-    zds = zarr.open_group(store, mode="r")
-    ds = open_dataset(store, engine="zarr", **kwargs)
-    tree_root = DataTree.from_dict({"/": ds})
-    for path in _iter_zarr_groups(zds):
-        try:
-            subgroup_ds = open_dataset(store, 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 _create_empty_netcdf_group(filename, group, mode, engine):
     ncDataset = _get_nc_dataset_class(engine)
 

xarray/datatree_/datatree/tests/conftest.py

@@ -1,7 +1,7 @@
 import pytest
 import xarray as xr
 
-from datatree import DataTree
+from xarray.datatree_.datatree import DataTree
 
 
 @pytest.fixture(scope="module")

xarray/datatree_/datatree/tests/test_dataset_api.py

@@ -1,8 +1,8 @@
 import numpy as np
 import xarray as xr
 
-from datatree import DataTree
-from datatree.testing import assert_equal
+from xarray.datatree_.datatree import DataTree
+from xarray.datatree_.datatree.testing import assert_equal
 
 
 class TestDSMethodInheritance:

xarray/datatree_/datatree/tests/test_datatree.py

@@ -6,8 +6,8 @@
 import xarray.testing as xrt
 from xarray.tests import create_test_data, source_ndarray
 
-import datatree.testing as dtt
-from datatree import DataTree, NotFoundInTreeError
+import xarray.datatree_.datatree.testing as dtt
+from xarray.datatree_.datatree import DataTree, NotFoundInTreeError
 
 
 class TestTreeCreation: