From a2a046a9c166ee3270427ccb8e712fffa8d83df2 Mon Sep 17 00:00:00 2001
From: Simon Heybrock <simon.heybrock@ess.eu>
Date: Thu, 10 Oct 2024 15:58:07 +0200
Subject: [PATCH] Add __delitem__

In line with how __getitem__ and __setitem__ work this cuts off the
branch "held" by the key, as well as removing the node data.

As discussed elsewhere the method naming is unusual, but if we intend to
change this it should be done elsewhere, after more thorough thinking
about the whole set of methods.
---
 src/cyclebane/graph.py | 18 ++++++++++++++
 tests/graph_test.py    | 55 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/src/cyclebane/graph.py b/src/cyclebane/graph.py
index 682e9c6..f2cfac7 100644
--- a/src/cyclebane/graph.py
+++ b/src/cyclebane/graph.py
@@ -401,6 +401,24 @@ def __getitem__(self, key: Hashable | slice) -> Graph:
             node_values=self._node_values.get_columns(keep_values),
         )
 
+    def __delitem__(self, key: Hashable | slice) -> None:
+        """
+        Delete the branch of the graph rooted at the given node.
+        """
+        if isinstance(key, slice):
+            raise NotImplementedError('Only single nodes are supported ')
+        key = self._from_orig_key(key)
+        if isinstance(key, MappedNode):
+            # Not clear what to do in this case, as it would leave a lot of MappedNodes
+            # without a source that could provide data.
+            raise ValueError('Cannot delete mapped node.')
+        graph = _remove_ancestors(self.graph, key)
+        graph.nodes[key].clear()
+        mapped = {node.name for node in graph if isinstance(node, MappedNode)}
+        keep_values = [key for key in self._node_values.keys() if key in mapped]
+        self._node_values = self._node_values.get_columns(keep_values)
+        self.graph = graph
+
     def __setitem__(self, branch: Hashable | slice, other: Graph) -> None:
         """
         Set a new branch in place of the given branch.
diff --git a/tests/graph_test.py b/tests/graph_test.py
index d036679..b9a3d27 100644
--- a/tests/graph_test.py
+++ b/tests/graph_test.py
@@ -558,6 +558,61 @@ def test_axis_in_reduce_refers_to_node_axis_not_graph_axis() -> None:
     assert all(n.index.axes == ('y',) for n in d_nodes)
 
 
+def test_delitem_removes_ancestors_and_data_but_keeps_node() -> None:
+    g = nx.DiGraph()
+    g.add_edge('a', 'b')
+    g.add_edge('b', 'c')
+    g.add_node('b', value='x')
+
+    graph = cb.Graph(g)
+    del graph['b']
+    result = graph.to_networkx()
+    assert list(result.nodes) == ['b', 'c']
+    assert result.nodes['b'] == {}
+
+
+def test_delitem_preserves_ancestors_with_other_path_to_graph() -> None:
+    g = nx.DiGraph()
+    g.add_edge('a', 'b')
+    g.add_edge('b', 'c')
+    g.add_node('b', value='x')
+    # a has another path to c so it will we kept
+    g.add_edge('a', 'c')
+
+    graph = cb.Graph(g)
+    del graph['b']
+    result = graph.to_networkx()
+    assert list(result.nodes) == ['a', 'b', 'c']
+    assert result.nodes['b'] == {}
+    assert not result.has_edge('a', 'b')
+
+
+def test_delitem_raises_when_removing_mapped_node() -> None:
+    g = nx.DiGraph()
+    g.add_edge('a', 'b')
+    g.add_edge('b', 'c')
+    g.add_node('b', value='x')
+
+    graph = cb.Graph(g)
+    mapped = graph.map({'a': [1, 2, 3]})
+    with pytest.raises(ValueError, match="Cannot delete mapped node."):
+        del mapped['b']
+
+
+def test_delitem_can_remove_reduced_node_depending_on_mapped_nodes() -> None:
+    g = nx.DiGraph()
+    g.add_edge('a', 'b')
+    g.add_edge('x', 'y')  # unrelated
+    graph = cb.Graph(g)
+    mapped = graph.map({'a': [1, 2, 3]}).reduce('b', name='c')
+
+    del mapped['c']
+
+    result = mapped.to_networkx()
+    assert set(result.nodes) == {'x', 'y', 'c'}
+    assert result.nodes['c'] == {}
+
+
 def test_setitem_raises_TypeError_if_given_networkx_graph() -> None:
     g = nx.DiGraph()
     g.add_edge('a', 'b')