Return (src, dst) in original order, and allow keys + data

This commit is contained in:
Bob Lantz
2014-10-28 20:43:10 -07:00
parent eab4ea3fb9
commit ba8ea8f0cc
+37 -25
View File
@@ -53,24 +53,28 @@ class MultiGraph( object ):
data: return list of ( node, attrs)""" data: return list of ( node, attrs)"""
return self.node.items() if data else self.node.keys() return self.node.items() if data else self.node.keys()
def edges_iter( self, data=False, key=False): def edges_iter( self, data=False, keys=False ):
"Iterator: return graph edges" "Iterator: return graph edges"
for src, entry in self.edge.iteritems(): for src, entry in self.edge.iteritems():
for dst, keys in entry.iteritems(): for dst, keys in entry.iteritems():
if [ src, dst ] != sorted( [ src, dst ] ): if src > dst:
# Skip duplicate edges # Skip duplicate edges
continue continue
for k, attrs in keys.iteritems(): for k, attrs in keys.iteritems():
if data: if data:
yield( ( src, dst, attrs ) ) if keys:
elif key: yield( src, dst, k, attrs )
yield( ( src, dst, k ) ) else:
yield( src, dst, attrs )
else: else:
yield( ( src, dst ) ) if keys:
yield( src, dst, k )
else:
yield( src, dst )
def edges( self, data=False, key=False ): def edges( self, data=False, keys=False ):
"Return list of graph edges" "Return list of graph edges"
return list( self.edges_iter( data=data, key=key ) ) return list( self.edges_iter( data=data, keys=key ) )
def __getitem__( self, node ): def __getitem__( self, node ):
@@ -174,24 +178,32 @@ class Topo( object ):
returns: list of hosts""" returns: list of hosts"""
return [ n for n in self.nodes( sort ) if not self.isSwitch( n ) ] return [ n for n in self.nodes( sort ) if not self.isSwitch( n ) ]
def links( self, sort=True, withInfo=False, withKeys=False ): def iterLinks( self, withKeys=False, withInfo=False ):
"""Return links. """Return links (iterator)
sort: sort links alphabetically withKeys: return link keys
withKeys: return key in tuple withInfo: return link info
returns: links list of ( src, dst [,key ] )""" returns: list of ( src, dst [,key, info ] )"""
if not sort: for src, dst, key, info in self.g.edges_iter( data=True, keys=True ):
return self.g.edges( data=withInfo, key=withKeys ) node1, node2 = info[ 'node1' ], info[ 'node2' ]
else: if withKeys:
if withInfo: if withInfo:
links = [ tuple( self.sorted( ( s, d ) ) + [ info ] ) yield( node1, node2, key, info )
for s, d, info in self.g.edges( data=True ) ] else:
elif withKeys: yield( node1, node2, key )
links = [ tuple( self.sorted( ( s, d ) ) + [ key ] )
for s, d, key in self.g.edges( key=True ) ]
else: else:
links = [ tuple ( self.sorted( e ) ) if withInfo:
for e in self.g.edges() ] yield( node1, node2, info )
return self.sorted( links ) else:
yield( node1, node2 )
def links( self, sort=False, withKeys=False, withInfo=False ):
"""Return links
sort: sort links alphabetically, preserving (src, dst) order
withKeys: return link keys
withInfo: return link info
returns: list of ( src, dst [,key, info ] )"""
links = list( self.iterLinks( withKeys, withInfo ) )
return links if not sort else self.sorted( links )
# This legacy port management mechanism is clunky and will probably # This legacy port management mechanism is clunky and will probably
# be removed at some point. # be removed at some point.