Fix DataTree.coords.__setitem__ by adding DataTreeCoordinates class

xarray/core/coordinates.py

@@ -36,6 +36,7 @@
     from xarray.core.common import DataWithCoords
     from xarray.core.dataarray import DataArray
     from xarray.core.dataset import Dataset
+    from xarray.core.datatree import DataTree
 
 # Used as the key corresponding to a DataArray's variable when converting
 # arbitrary DataArray objects to datasets
@@ -197,12 +198,12 @@ class Coordinates(AbstractCoordinates):
 
     Coordinates are either:
 
-    - returned via the :py:attr:`Dataset.coords` and :py:attr:`DataArray.coords`
-      properties
+    - returned via the :py:attr:`Dataset.coords`, :py:attr:`DataArray.coords`,
+      and :py:attr:`DataTree.coords` properties,
     - built from Pandas or other index objects
-      (e.g., :py:meth:`Coordinates.from_pandas_multiindex`)
+      (e.g., :py:meth:`Coordinates.from_pandas_multiindex`),
     - built directly from coordinate data and Xarray ``Index`` objects (beware that
-      no consistency check is done on those inputs)
+      no consistency check is done on those inputs),
 
     Parameters
     ----------
@@ -796,6 +797,126 @@ def _ipython_key_completions_(self):
         ]
 
 
+class DataTreeCoordinates(DatasetCoordinates):
+    """
+    Dictionary like container for coordinates of a DataTree node (variables + indexes).
+    """
+
+    # TODO: This collection can be passed directly to the :py:class:`~xarray.Dataset`
+    # and :py:class:`~xarray.DataArray` constructors via their `coords` argument.
+    # This will add both the coordinates variables and their index.
+
+    # TODO: This only needs to be a separate class from `DatasetCoordinates` because DataTree nodes store their variables differently
+    # internally than how Datasets do, see https://github.com/pydata/xarray/issues/9203.
+
+    # TODO should inherited coordinates be here? It would be very hard to allow updating them...
+    # But actually maybe the ChainMap approach would make this work okay??
+
+    _data: DataTree
+
+    __slots__ = ("_data",)
+
+    def __init__(self, datatree: DataTree):
+        self._data = datatree
+
+    @property
+    def _names(self) -> set[Hashable]:
+        return set(self._data._coord_variables)
+
+    @property
+    def dims(self) -> Frozen[Hashable, int]:
+        # TODO is there a potential bug here? What happens if a dim is only present on data variables?
+        return self._data.dims
+
+    @property
+    def dtypes(self) -> Frozen[Hashable, np.dtype]:
+        """Mapping from coordinate names to dtypes.
+
+        Cannot be modified directly, but is updated when adding new variables.
+
+        See Also
+        --------
+        Dataset.dtypes
+        """
+        return Frozen({n: v.dtype for n, v in self._data._coord_variables.items()})
+
+    @property
+    def variables(self) -> Mapping[Hashable, Variable]:
+        variables = self._data._data_variables | self._data._coord_variables
+        return Frozen({k: v for k, v in variables.items() if k in self._names})
+
+    def __getitem__(self, key: Hashable) -> DataArray:
+        if key in self._data._data_variables or key in self._data.children:
+            raise KeyError(key)
+        return self._data[key]
+
+    def to_dataset(self) -> Dataset:
+        """Convert these coordinates into a new Dataset"""
+        return self._data.to_dataset()._copy_listed(self._names)
+
+    def _update_coords(
+        self, coords: dict[Hashable, Variable], indexes: Mapping[Any, Index]
+    ) -> None:
+
+        # TODO I don't know how to update coordinates that live in parent nodes
+        # TODO We would have to find the correct node and update `._node_coord_variables`
+
+        coord_variables = self._data._coord_variables.copy()
+        coord_variables.update(coords)
+
+        # check for inconsistent state *before* modifying anything in-place
+        variables = coord_variables | self._data._data_variables.copy()
+        # TODO is there a subtlety here with rebuild_dims?
+        dims = calculate_dimensions(variables)
+        new_coord_names = set(coords)
+        for dim, size in dims.items():
+            if dim in variables:
+                new_coord_names.add(dim)
+
+        # TODO we need to upgrade these variables to coord variables somehow
+        # coord_variables.update(new_coord_names)
+
+        self._data._coord_variables = coord_variables
+        self._data._dims = dims
+
+        # TODO(shoyer): once ._indexes is always populated by a dict, modify
+        # it to update inplace instead.
+        original_indexes = dict(self._data.xindexes)
+        original_indexes.update(indexes)
+        self._data._indexes = original_indexes
+
+    def _drop_coords(self, coord_names):
+        # should drop indexed coordinates only
+        for name in coord_names:
+            del self._data._coord_variables[name]
+            del self._data._indexes[name]
+        # self._data._coord_names.difference_update(coord_names)
+
+    def _drop_indexed_coords(self, coords_to_drop: set[Hashable]) -> None:
+        assert self._data.xindexes is not None
+        new_coords = drop_indexed_coords(coords_to_drop, self)
+        for name in self._data._coord_names - new_coords._names:
+            del self._data._coord_variables[name]
+        self._data._indexes = dict(new_coords.xindexes)
+        # self._data._coord_names.intersection_update(new_coords._names)
+
+    def __delitem__(self, key: Hashable) -> None:
+        if key in self:
+            del self._data[key]
+        else:
+            raise KeyError(
+                f"{key!r} is not in coordinate variables {tuple(self.keys())}"
+            )
+
+    def _ipython_key_completions_(self):
+        """Provide method for the key-autocompletions in IPython."""
+        return [
+            key
+            for key in self._data._ipython_key_completions_()
+            if key not in self._data.data_vars
+        ]
+
+
 class DataArrayCoordinates(Coordinates, Generic[T_DataArray]):
     """Dictionary like container for DataArray coordinates (variables + indexes).
 

xarray/core/datatree.py

@@ -16,7 +16,7 @@
 from xarray.core import utils
 from xarray.core.alignment import align
 from xarray.core.common import TreeAttrAccessMixin
-from xarray.core.coordinates import DatasetCoordinates
+from xarray.core.coordinates import DataTreeCoordinates
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset, DataVariables
 from xarray.core.datatree_mapping import (
@@ -1148,11 +1148,11 @@ def xindexes(self) -> Indexes[Index]:
         )
 
     @property
-    def coords(self) -> DatasetCoordinates:
+    def coords(self) -> DataTreeCoordinates:
         """Dictionary of xarray.DataArray objects corresponding to coordinate
         variables
         """
-        return DatasetCoordinates(self.to_dataset())
+        return DataTreeCoordinates(self)
 
     @property
     def data_vars(self) -> DataVariables:

xarray/tests/test_datatree.py

@@ -7,11 +7,13 @@
 import pytest
 
 import xarray as xr
+from xarray import Dataset
+from xarray.core.coordinates import DataTreeCoordinates
 from xarray.core.datatree import DataTree
 from xarray.core.datatree_ops import _MAPPED_DOCSTRING_ADDENDUM, insert_doc_addendum
 from xarray.core.treenode import NotFoundInTreeError
 from xarray.testing import assert_equal, assert_identical
-from xarray.tests import create_test_data, source_ndarray
+from xarray.tests import assert_array_equal, create_test_data, source_ndarray
 
 
 class TestTreeCreation:
@@ -525,6 +527,122 @@
         assert_identical(results.to_dataset(), expected)
 
 
+class TestCoordsInterface:
+    def test_properties(self):
+        # use int64 for repr consistency on windows
+        ds = Dataset(
+            data_vars={
+                "foo": (["x", "y"], np.random.randn(2, 3)),
+            },
+            coords={
+                "x": ("x", np.array([-1, -2], "int64")),
+                "y": ("y", np.array([0, 1, 2], "int64")),
+                "a": ("x", np.array([4, 5], "int64")),
+                "b": np.int64(-10),
+            },
+        )
+        dt = DataTree(data=ds)
+        dt["child"] = DataTree()
+
+        coords = dt.coords
+        assert isinstance(coords, DataTreeCoordinates)
+
+        # len
+        assert len(coords) == 4
+
+        # iter
+        assert list(coords) == ["x", "y", "a", "b"]
+
+        assert_identical(coords["x"].variable, dt["x"].variable)
+        assert_identical(coords["y"].variable, dt["y"].variable)
+
+        assert "x" in coords
+        assert "a" in coords
+        assert 0 not in coords
+        assert "foo" not in coords
+        assert "child" not in coords
+
+        with pytest.raises(KeyError):
+            coords["foo"]
+
+        # TODO this currently raises a ValueError instead of a KeyError
+        # with pytest.raises(KeyError):
+        #     coords[0]
+
+        # repr
+        expected = dedent(
+            """\
+        Coordinates:
+          * x        (x) int64 16B -1 -2
+          * y        (y) int64 24B 0 1 2
+            a        (x) int64 16B 4 5
+            b        int64 8B -10"""
+        )
+        actual = repr(coords)
+        assert expected == actual
+
+        # dims
+        assert coords.sizes == {"x": 2, "y": 3}
+
+        # dtypes
+        assert coords.dtypes == {
+            "x": np.dtype("int64"),
+            "y": np.dtype("int64"),
+            "a": np.dtype("int64"),
+            "b": np.dtype("int64"),
+        }
+
+    def test_modify(self):
+        ds = Dataset(
+            data_vars={
+                "foo": (["x", "y"], np.random.randn(2, 3)),
+            },
+            coords={
+                "x": ("x", np.array([-1, -2], "int64")),
+                "y": ("y", np.array([0, 1, 2], "int64")),
+                "a": ("x", np.array([4, 5], "int64")),
+                "b": np.int64(-10),
+            },
+        )
+        dt = DataTree(data=ds)
+        dt["child"] = DataTree()
+
+        actual = dt.copy(deep=True)
+        actual.coords["x"] = ("x", ["a", "b"])
+        assert_array_equal(actual["x"], ["a", "b"])
+
+        actual = dt.copy(deep=True)
+        actual.coords["z"] = ("z", ["a", "b"])
+        assert_array_equal(actual["z"], ["a", "b"])
+
+        actual = dt.copy(deep=True)
+        with pytest.raises(ValueError, match=r"conflicting dimension sizes"):
+            actual.coords["x"] = ("x", [-1])
+        assert_identical(actual, dt)  # should not be modified
+
+        actual = dt.copy()
+        del actual.coords["b"]
+        expected = dt.reset_coords("b", drop=True)
+        assert_identical(expected, actual)
+
+        with pytest.raises(KeyError):
+            del dt.coords["not_found"]
+
+        with pytest.raises(KeyError):
+            del dt.coords["foo"]
+
+        actual = dt.copy(deep=True)
+        actual.coords.update({"c": 11})
+        expected = dt.merge({"c": 11}).set_coords("c")
+        assert_identical(expected, actual)
+
+        # regression test for GH3746
+        del actual.coords["x"]
+        assert "x" not in actual.xindexes
+
+    # TODO test with coordinate inheritance too...
+
+
 class TestDictionaryInterface: ...