Skip to content

Commit

Permalink
Add __delitem__
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
SimonHeybrock committed Oct 10, 2024
1 parent ee09d24 commit a2a046a
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 0 deletions.
18 changes: 18 additions & 0 deletions src/cyclebane/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
55 changes: 55 additions & 0 deletions tests/graph_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down

0 comments on commit a2a046a

Please sign in to comment.