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): ...@@ -7,6 +7,7 @@ class DirectedGraph(object):
""" """
def __init__(self): def __init__(self):
self._edges = {} self._edges = {}
self._edge_numbers = {}
def add_node(self, p_id): def add_node(self, p_id):
""" Adds a node to the graph. """ """ Adds a node to the graph. """
...@@ -17,10 +18,12 @@ class DirectedGraph(object): ...@@ -17,10 +18,12 @@ class DirectedGraph(object):
""" Returns true iff the graph has the given node. """ """ Returns true iff the graph has the given node. """
return p_id in self._edges 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 Adds an edge to the graph. The nodes will be added if they don't
exist. exist.
The p_id is the id of the edge, if the client wishes to maintain this.
""" """
if not self.has_node(p_from): if not self.has_node(p_from):
self.add_node(p_from) self.add_node(p_from)
...@@ -29,6 +32,8 @@ class DirectedGraph(object): ...@@ -29,6 +32,8 @@ class DirectedGraph(object):
self.add_node(p_to) self.add_node(p_to)
self._edges[p_from].add(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): def has_path(self, p_from, p_to):
""" """
...@@ -116,6 +121,12 @@ class DirectedGraph(object): ...@@ -116,6 +121,12 @@ class DirectedGraph(object):
""" Returns True when the graph has the given edge. """ """ Returns True when the graph has the given edge. """
return p_from in self._edges and p_to in self._edges[p_from] 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): def remove_edge(self, p_from, p_to, remove_unconnected_nodes=True):
""" """
Removes an edge from the graph. Removes an edge from the graph.
...@@ -126,6 +137,11 @@ class DirectedGraph(object): ...@@ -126,6 +137,11 @@ class DirectedGraph(object):
if self.has_edge(p_from, p_to): if self.has_edge(p_from, p_to):
self._edges[p_from].remove(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 remove_unconnected_nodes:
if self.is_isolated(p_from): if self.is_isolated(p_from):
self.remove_node(p_from) self.remove_node(p_from)
......
...@@ -6,7 +6,7 @@ class GraphTest(unittest.TestCase): ...@@ -6,7 +6,7 @@ class GraphTest(unittest.TestCase):
def setUp(self): def setUp(self):
self.graph = Graph.DirectedGraph() 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(2, 4)
self.graph.add_edge(4, 3) self.graph.add_edge(4, 3)
self.graph.add_edge(4, 6) self.graph.add_edge(4, 6)
...@@ -58,6 +58,7 @@ class GraphTest(unittest.TestCase): ...@@ -58,6 +58,7 @@ class GraphTest(unittest.TestCase):
self.assertFalse(self.graph.has_path(1, 4)) self.assertFalse(self.graph.has_path(1, 4))
self.assertTrue(self.graph.has_path(2, 4)) self.assertTrue(self.graph.has_path(2, 4))
self.assertFalse(self.graph.has_edge_id(1))
def test_remove_edge2(self): def test_remove_edge2(self):
self.graph.remove_edge(3, 5, True) 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