the big merge

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

diff --git a/src/nepi/util/netgraph.py b/src/nepi/util/netgraph.py
index 7a681ec..96fb384 100644 (file)
--- a/src/nepi/util/netgraph.py
+++ b/src/nepi/util/netgraph.py
@@ -16,11 +16,17 @@
 #
 # Author: Alina Quereilhac <alina.quereilhac@inria.fr>
 
-import ipaddr
 import networkx
 import math
 import random
 
+from six import next, PY2, PY3
+if PY2:
+    import ipaddr
+else:
+    import ipaddress
+
+
 class TopologyType:
     LINEAR = "linear"
     LADDER = "ladder"
@@ -196,7 +202,9 @@ class NetGraph(object):
         return self.topology.node[nid].get(name)
 
     def node_annotations(self, nid):
-        return self.topology.node[nid].keys()
+        retcod = self.topology.node[nid].keys()
+        if PY3: retcod = list(retcod)
+        return retcod
     
     def del_node_annotation(self, nid, name):
         del self.topology.node[nid][name]
@@ -224,7 +232,9 @@ class NetGraph(object):
         return self.topology.edge[nid1][nid2].get(name)
  
     def edge_annotations(self, nid1, nid2):
-        return self.topology.edge[nid1][nid2].keys()
+        retcod = self.topology.edge[nid1][nid2].keys()
+        if PY3: retcod = list(retcod)
+        return retcod
     
     def del_edge_annotation(self, nid1, nid2, name):
         del self.topology.edge[nid1][nid2][name]
@@ -250,10 +260,10 @@ class NetGraph(object):
         # Assign IP addresses to host
         netblock = "%s/%d" % (network, prefix)
         if version == 4:
-            net = ipaddr.IPv4Network(netblock)
+            net = ipaddr.IPv4Network(netblock) if PY2 else ipaddress.ip_network(netblock)
             new_prefix = 30
         elif version == 6:
-            net = ipaddr.IPv6Network(netblock)
+            net = ipaddr.IPv6Network(netblock) if PY2 else ipaddress.ip_network(netblock)
             new_prefix = 30
         else:
             raise RuntimeError("Invalid IP version %d" % version)
@@ -269,15 +279,15 @@ class NetGraph(object):
             #### Compute subnets for each link
             
             # get a subnet of base_add with prefix /30
-            subnet = sub_itr.next()
+            subnet = next(sub_itr)
             mask = subnet.netmask.exploded
             network = subnet.network.exploded
             prefixlen = subnet.prefixlen
 
             # get host addresses in that subnet
             i = subnet.iterhosts()
-            addr1 = i.next()
-            addr2 = i.next()
+            addr1 = next(i)
+            addr2 = next(i)
 
             ip1 = addr1.exploded
             ip2 = addr2.exploded