Migrate datatree.py module into xarray.core. (#8789)

* Migrate datatree.py module into xarray.core.

* Add correct PR reference to whats-new.rst.

* Revert to using Union in datatree.py.

* Catch remaining unfixed import path.

* Fix easier mypy annotations in datatree.py and test_datatree.py.

* Straggling mypy change in datatree.py.

* datatree.py comment clean-up.

* More mypy corrections in datatree.py and test_datatree.py.

* Removes unnecessary dict wrapper.

* DAS-2062: renames as_array -> to_dataarray

* DAS-2062: Updates doc string for Datatree.to_zarr

Accurately reflects the default value now.

* DAS-2062: reverts what-new.rst

for "breaking changes that aren't breaking yet but will be, but only relevant
for previous users of another package"

* DAS-2062: clarify wording in comment.

* Change Datatree.to_dataarray to call correctly

We updated the name but not the function.

* Clarify DataTree's names are still strings now.

DAS-2062

* DAS-2062: Cast k explicitly to str for typing.

So this is where we are moving forward with the assumption that DataTree nodes
are alway named with a string. In this section of `update` even though we know
the key is a str, mypy refuses.

I chose explicit recast over mypy ignores, tell me why that's wrong?

* Ignore mypy errors for DataTree.ds assignment.

* Fix DataTree.update type hint.

* Final mypy issue - ignore DataTree.get override.

* Update contributors in whats-new.rst

* Fix GitHub handle.

---------

Co-authored-by: Matt Savoie <[email protected]>
Co-authored-by: Matt Savoie <[email protected]>
Co-authored-by: Deepak Cherian <[email protected]>

Author: 3 people

doc/whats-new.rst

@@ -82,6 +82,9 @@ Internal Changes
 - Migrates ``treenode`` functionality into ``xarray/core`` (:pull:`8757`)
   By `Matt Savoie <https://github.com/flamingbear>`_ and `Tom Nicholas
   <https://github.com/TomNicholas>`_.
+- Migrates ``datatree`` functionality into ``xarray/core``. (:pull: `8789`)
+  By `Owen Littlejohns <https://github.com/owenlittlejohns>`_, `Matt Savoie
+  <https://github.com/flamingbear>`_ and `Tom Nicholas <https://github.com/TomNicholas>`_.
 
 
 .. _whats-new.2024.02.0:

xarray/backends/api.py

@@ -69,7 +69,7 @@
     T_NetcdfTypes = Literal[
         "NETCDF4", "NETCDF4_CLASSIC", "NETCDF3_64BIT", "NETCDF3_CLASSIC"
     ]
-    from xarray.datatree_.datatree import DataTree
+    from xarray.core.datatree import DataTree
 
 DATAARRAY_NAME = "__xarray_dataarray_name__"
 DATAARRAY_VARIABLE = "__xarray_dataarray_variable__"

xarray/backends/common.py

@@ -23,8 +23,8 @@
     from netCDF4 import Dataset as ncDataset
 
     from xarray.core.dataset import Dataset
+    from xarray.core.datatree import DataTree
     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__)
@@ -137,8 +137,8 @@ def _open_datatree_netcdf(
     **kwargs,
 ) -> DataTree:
     from xarray.backends.api import open_dataset
+    from xarray.core.datatree import DataTree
     from xarray.core.treenode import NodePath
-    from xarray.datatree_.datatree import DataTree
 
     ds = open_dataset(filename_or_obj, **kwargs)
     tree_root = DataTree.from_dict({"/": ds})

xarray/backends/h5netcdf_.py

@@ -40,7 +40,7 @@
 
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
-    from xarray.datatree_.datatree import DataTree
+    from xarray.core.datatree import DataTree
 
 
 class H5NetCDFArrayWrapper(BaseNetCDF4Array):

xarray/backends/netCDF4_.py

@@ -45,7 +45,7 @@
 
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
-    from xarray.datatree_.datatree import DataTree
+    from xarray.core.datatree import DataTree
 
 # This lookup table maps from dtype.byteorder to a readable endian
 # string used by netCDF4.

xarray/backends/zarr.py

@@ -34,7 +34,7 @@
 
     from xarray.backends.common import AbstractDataStore
     from xarray.core.dataset import Dataset
-    from xarray.datatree_.datatree import DataTree
+    from xarray.core.datatree import DataTree
 
 
 # need some special secret attributes to tell us the dimensions
@@ -1048,8 +1048,8 @@ def open_datatree(
         import zarr
 
         from xarray.backends.api import open_dataset
+        from xarray.core.datatree import DataTree
         from xarray.core.treenode import NodePath
-        from xarray.datatree_.datatree import DataTree
 
         zds = zarr.open_group(filename_or_obj, mode="r")
         ds = open_dataset(filename_or_obj, engine="zarr", **kwargs)