Bug fixing the netgraph persistence

[nepi.git] / src / nepi / util / netgraph.py

diff --git a/src/nepi/util/netgraph.py b/src/nepi/util/netgraph.py
index c04f94d..4b8c06c 100644 (file)
--- a/src/nepi/util/netgraph.py
+++ b/src/nepi/util/netgraph.py
@@ -70,10 +70,18 @@ class NetGraph(object):
             :param version: IP version for IP address assignment.
             :type version: int
 
-
             :param assign_st: Select source and target nodes on the graph.
             :type assign_st: bool
 
+           :param sources_targets: dictionary with the list of sources (key =
+            "sources") and list of targets (key = "targets") if defined, ignore
+            assign_st
+           :type sources_targets: dictionary of lists
+
+           :param leaf_source: if True, random sources will be selected only 
+            from leaf nodes.
+           :type leaf_source: bool
+
         NOTE: Only point-to-point like network topologies are supported for now.
                 (Wireless and Ethernet networks were several nodes share the same
                 edge (hyperedge) can not be modeled for the moment).
@@ -100,9 +108,13 @@ class NetGraph(object):
             self.assign_p2p_ips(network = network, prefix = prefix, 
                     version = version)
 
-        if kwargs.get("assign_st"):
+       sources_targets = kwargs.get("sources_targets")
+       if sources_targets:
+            [self.set_source(n) for n in sources_targets["sources"]]
+           [self.set_target(n) for n in sources_targets["targets"]]
+       elif kwargs.get("assign_st"):
             self.select_target_zero()
-            self.select_random_leaf_source()
+            self.select_random_source(is_leaf = kwargs.get("leaf_source"))
 
     @property
     def topology(self):
@@ -152,24 +164,13 @@ class NetGraph(object):
                     prev = c
                     c += 1
 
-        # node ids are int, make them str
-        g = networkx.Graph()
-        g.add_nodes_from(map(lambda nid: str(nid), graph.nodes()))
-        g.add_edges_from(map(lambda t: (str(t[0]), str(t[1])), 
-            graph.edges()))
-
-        return g
+        return graph
 
     def add_node(self, nid):
-        nid = str(nid)
-
         if nid not in self.topology: 
             self.topology.add_node(nid)
 
     def add_edge(self, nid1, nid2):
-        nid1 = str(nid1)
-        nid2 = str(nid2)
-
         self.add_node(nid1)
         self.add_node( nid2)
 
@@ -315,27 +316,28 @@ class NetGraph(object):
                 if self.topology.node[nid].get("source")]
 
     def select_target_zero(self):
-        """ Marks the node 0 as target
+        """ Mark the node 0 as target
         """
-        self.set_target("0")
+        nid = 0 if 0 in self.topology.nodes() else "0"
+        self.set_target(nid)
 
-    def select_random_leaf_source(self):
-        """  Marks a random leaf node as source. 
+    def select_random_source(self, **kwargs):
+        """  Mark a random node as source. 
         """
 
         # The ladder is a special case because is not symmetric.
         if self.topo_type == TopologyType.LADDER:
             total_nodes = self.order/2
-            leaf1 = str(total_nodes)
-            leaf2 = str(total_nodes - 1)
+            leaf1 = total_nodes
+            leaf2 = total_nodes - 1
             leaves = [leaf1, leaf2]
             source = leaves.pop(random.randint(0, len(leaves) - 1))
         else:
             # options must not be already sources or targets
             options = [ k for k,v in self.topology.degree().iteritems() \
-                    if v == 1 and not self.topology.node[k].get("source") \
+                    if (not kwargs.get("is_leaf") or v == 1)  \
+                        and not self.topology.node[k].get("source") \
                         and not self.topology.node[k].get("target")]
-
             source = options.pop(random.randint(0, len(options) - 1))
         
         self.set_source(source)