From 9b5056761b3bc1720a36326e3d3ad557b6bb94d6 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Sat, 7 Sep 2024 23:36:21 -0400 Subject: [PATCH 1/7] test --- xarray/tests/test_datatree.py | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 1c47726b470..f9fba107d0a 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -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 @@ -525,7 +526,28 @@ 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()}) + + # test delete children + del dt["c"] + assert dt.children == {} + assert set(dt.variables) == {"x", "z", "a"} + + # test delete variables + del dt["a"] + assert set(dt.coords) == {"x", "z"} + + # test delete coordinates + del dt["z"] + assert set(dt.coords) == {"x"} + + # test delete indexed coordinates + del dt["x"] + assert dt.variables == {} + assert dt.coords == {} + assert dt.indexes == {} class TestTreeFromDict: From c466f8db0fa2ed7b89d2e6791e931f005a70a14a Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Sat, 7 Sep 2024 23:36:55 -0400 Subject: [PATCH 2/7] improve error message --- xarray/core/treenode.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 9dfd346508a..8f7d5d7fc1f 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -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(f"Child node {key} not found") def same_tree(self, other: Tree) -> bool: """True if other node is in the same tree as this node.""" From 0397eca3065397ae4c73aea3519058939e5ff38b Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Sat, 7 Sep 2024 23:37:04 -0400 Subject: [PATCH 3/7] implement delitem --- xarray/core/datatree.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index cd0c647f491..00112755dba 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -910,6 +910,26 @@ 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( + f"Cannot delete {key} as it was not found on this datatree node. Must be one of {list(self)}" + ) + @overload def update(self, other: Dataset) -> None: ... From 85bb221ec930a06956bce0e26f984df44ce6e52e Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Sat, 7 Sep 2024 23:38:42 -0400 Subject: [PATCH 4/7] test KeyError --- xarray/core/datatree.py | 2 +- xarray/tests/test_datatree.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index 00112755dba..8f56daafc86 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -927,7 +927,7 @@ def __delitem__(self, key: str) -> None: else: raise KeyError( - f"Cannot delete {key} as it was not found on this datatree node. Must be one of {list(self)}" + f"Cannot delete key '{key}' as it was not found on this datatree node. Must be one of {list(self)}" ) @overload diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index f9fba107d0a..077e629186a 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -549,6 +549,9 @@ def test_delitem(): assert dt.coords == {} assert dt.indexes == {} + with pytest.raises(KeyError, match="Cannot delete key 'foo' as it was not found"): + del dt["foo"] + class TestTreeFromDict: def test_data_in_root(self): From b6738b5d1d172022a358b375691ced308fe06fc4 Mon Sep 17 00:00:00 2001 From: Tom Nicholas Date: Sun, 8 Sep 2024 14:22:07 -0600 Subject: [PATCH 5/7] Remove special KeyError message Co-authored-by: Stephan Hoyer --- xarray/core/datatree.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index 8f56daafc86..3483a0e5755 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -926,9 +926,7 @@ def __delitem__(self, key: str) -> None: self._node_dims = calculate_dimensions(self.variables) else: - raise KeyError( - f"Cannot delete key '{key}' as it was not found on this datatree node. Must be one of {list(self)}" - ) + raise KeyError(key) @overload def update(self, other: Dataset) -> None: ... From 3c3b93f2423c67822a5d2ce4572c9f140fe80935 Mon Sep 17 00:00:00 2001 From: Tom Nicholas Date: Sun, 8 Sep 2024 14:22:28 -0600 Subject: [PATCH 6/7] remove another special KeyError message Co-authored-by: Stephan Hoyer --- xarray/core/treenode.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 8f7d5d7fc1f..84ce392ad32 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -566,7 +566,7 @@ def __delitem__(self: Tree, key: str) -> None: del self._children[key] child.orphan() else: - raise KeyError(f"Child node {key} not found") + raise KeyError(key) def same_tree(self, other: Tree) -> bool: """True if other node is in the same tree as this node.""" From 3dbe493de571e43bb3e58aa96190010a3cc2b616 Mon Sep 17 00:00:00 2001 From: TomNicholas Date: Mon, 9 Sep 2024 11:08:24 -0400 Subject: [PATCH 7/7] fix pytest raises expected error message --- xarray/tests/test_datatree.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index 077e629186a..a2ff043d02f 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -530,27 +530,35 @@ 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, match="Cannot delete key 'foo' as it was not found"): - del dt["foo"] + with pytest.raises(KeyError): + del dt["x"] class TestTreeFromDict: