Skip to content

Commit

Permalink
Implement DataTree.__delitem__ (#9453)
Browse files Browse the repository at this point in the history
* test

* improve error message

* implement delitem

* test KeyError

* Remove special KeyError message

Co-authored-by: Stephan Hoyer <shoyer@google.com>

* remove another special KeyError message

Co-authored-by: Stephan Hoyer <shoyer@google.com>

* fix pytest raises expected error message

---------

Co-authored-by: Stephan Hoyer <shoyer@google.com>
  • Loading branch information
TomNicholas and shoyer authored Sep 9, 2024
1 parent 0af1979 commit 45a0027
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 3 deletions.
18 changes: 18 additions & 0 deletions xarray/core/datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -910,6 +910,24 @@ def __setitem__(
else:
raise ValueError("Invalid format for key")

def __delitem__(self, key: str) -> None:
"""Remove a variable or child node from this datatree node."""
if key in self.children:
super().__delitem__(key)

elif key in self._node_coord_variables:
if key in self._node_indexes:
del self._node_indexes[key]
del self._node_coord_variables[key]
self._node_dims = calculate_dimensions(self.variables)

elif key in self._data_variables:
del self._data_variables[key]
self._node_dims = calculate_dimensions(self.variables)

else:
raise KeyError(key)

@overload
def update(self, other: Dataset) -> None: ...

Expand Down
4 changes: 2 additions & 2 deletions xarray/core/treenode.py
Original file line number Diff line number Diff line change
Expand Up @@ -559,14 +559,14 @@ def _set_item(
else:
current_node._set(name, item)

def __delitem__(self: Tree, key: str):
def __delitem__(self: Tree, key: str) -> None:
"""Remove a child node from this tree object."""
if key in self.children:
child = self._children[key]
del self._children[key]
child.orphan()
else:
raise KeyError("Cannot delete")
raise KeyError(key)

def same_tree(self, other: Tree) -> bool:
"""True if other node is in the same tree as this node."""
Expand Down
35 changes: 34 additions & 1 deletion xarray/tests/test_datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import pytest

import xarray as xr
from xarray import Dataset
from xarray.core.datatree import DataTree
from xarray.core.datatree_ops import _MAPPED_DOCSTRING_ADDENDUM, insert_doc_addendum
from xarray.core.treenode import NotFoundInTreeError
Expand Down Expand Up @@ -525,7 +526,39 @@ def test_setitem_dataarray_replace_existing_node(self):
assert_identical(results.to_dataset(), expected)


class TestDictionaryInterface: ...
def test_delitem():
ds = Dataset({"a": 0}, coords={"x": ("x", [1, 2]), "z": "a"})
dt = DataTree(ds, children={"c": DataTree()})

with pytest.raises(KeyError):
del dt["foo"]

# test delete children
del dt["c"]
assert dt.children == {}
assert set(dt.variables) == {"x", "z", "a"}
with pytest.raises(KeyError):
del dt["c"]

# test delete variables
del dt["a"]
assert set(dt.coords) == {"x", "z"}
with pytest.raises(KeyError):
del dt["a"]

# test delete coordinates
del dt["z"]
assert set(dt.coords) == {"x"}
with pytest.raises(KeyError):
del dt["z"]

# test delete indexed coordinates
del dt["x"]
assert dt.variables == {}
assert dt.coords == {}
assert dt.indexes == {}
with pytest.raises(KeyError):
del dt["x"]


class TestTreeFromDict:
Expand Down

0 comments on commit 45a0027

Please sign in to comment.