Commit 4b5e72b7 authored by Bram Schoenmakers's avatar Bram Schoenmakers

Also add the ability to store an identifier with an edge.

In fact, the id: and p: tags in a todo are edges, so this value has to
be stored somewhere.
parent e9e2e77d
......@@ -7,6 +7,7 @@ class DirectedGraph(object):
"""
def __init__(self):
self._edges = {}
self._edge_numbers = {}
def add_node(self, p_id):
""" Adds a node to the graph. """
......@@ -17,10 +18,12 @@ class DirectedGraph(object):
""" Returns true iff the graph has the given node. """
return p_id in self._edges
def add_edge(self, p_from, p_to):
def add_edge(self, p_from, p_to, p_id=None):
"""
Adds an edge to the graph. The nodes will be added if they don't
exist.
The p_id is the id of the edge, if the client wishes to maintain this.
"""
if not self.has_node(p_from):
self.add_node(p_from)
......@@ -29,6 +32,8 @@ class DirectedGraph(object):
self.add_node(p_to)
self._edges[p_from].add(p_to)
if p_id:
self._edge_numbers[p_id] = (p_from, p_to)
def has_path(self, p_from, p_to):
"""
......@@ -116,6 +121,12 @@ class DirectedGraph(object):
""" Returns True when the graph has the given edge. """
return p_from in self._edges and p_to in self._edges[p_from]
def has_edge_id(self, p_id):
"""
Returns True if the client registered an edge with the given id.
"""
return p_id in self._edge_numbers
def remove_edge(self, p_from, p_to, remove_unconnected_nodes=True):
"""
Removes an edge from the graph.
......@@ -126,6 +137,11 @@ class DirectedGraph(object):
if self.has_edge(p_from, p_to):
self._edges[p_from].remove(p_to)
for key, value in self._edge_numbers.iteritems():
if value == (p_from, p_to):
del self._edge_numbers[key]
break
if remove_unconnected_nodes:
if self.is_isolated(p_from):
self.remove_node(p_from)
......
......@@ -6,7 +6,7 @@ class GraphTest(unittest.TestCase):
def setUp(self):
self.graph = Graph.DirectedGraph()
self.graph.add_edge(1, 2)
self.graph.add_edge(1, 2, 1)
self.graph.add_edge(2, 4)
self.graph.add_edge(4, 3)
self.graph.add_edge(4, 6)
......@@ -58,6 +58,7 @@ class GraphTest(unittest.TestCase):
self.assertFalse(self.graph.has_path(1, 4))
self.assertTrue(self.graph.has_path(2, 4))
self.assertFalse(self.graph.has_edge_id(1))
def test_remove_edge2(self):
self.graph.remove_edge(3, 5, True)
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment