elements: Add exporting of DIF graphs

This adds the feature to allow to export a DIF graph to
PDF. Enrollments are shown as black, management flows as blue and data
transfer flows as red.

1 files changed, 112 insertions, 36 deletions

diff --git a/rumba/elements/experimentation.py b/rumba/elements/experimentation.py
index 60a35e1..e62103f 100644
--- a/rumba/elements/experimentation.py
+++ b/rumba/elements/experimentation.py
@@ -545,50 +545,126 @@ class Experiment(object):
         """
         try:
             import pydot
+
+            colors = ['red', 'green', 'blue', 'orange', 'yellow']
+            fcolors = ['black', 'black', 'white', 'black', 'black']
+
+            gvizg = pydot.Dot(graph_type='graph')
+
+            i = 0
+            try:
+                for node in self.nodes:
+                    g_node = pydot.Node(node.name,
+                                        label=node.name,
+                                        style="filled",
+                                        fillcolor=colors[i],
+                                        fontcolor=fcolors[i])
+
+                    gvizg.add_node(g_node)
+                    i += 1
+                    if i == len(colors):
+                        i = 0
+            except Exception as e:
+                logger.error('Failed to create pydot Node: ' + str(e))
+                return
+
+            try:
+                for dif in self.dif_ordering:
+                    if isinstance(dif, topology.ShimEthDIF):
+                        edge = pydot.Edge(dif.members[0].name,
+                                          dif.members[1].name,
+                                          label=dif.name,
+                                          color='black')
+                        gvizg.add_edge(edge)
+            except Exception as e:
+                logger.error('Failed to create pydot Edge: ' + str(e))
+                return
+
+            try:
+                gvizg.write_pdf(filename)
+                logger.info('Generated pdf of connectivity graph')
+            except Exception as e:
+                logger.error('Failed to write PDF: ' + str(e))
+
         except:
             logger.error('Warning: pydot module not installed, '
-                         'cannot produce DIF graphs images')
-            return
+                         'cannot produce graph images')
 
-        colors = ['red', 'green', 'blue', 'orange', 'yellow']
-        fcolors = ['black', 'black', 'white', 'black', 'black']
-
-        gvizg = pydot.Dot(graph_type='graph')
+    def export_dif_graph(self, filename, dif):
+        """
+        Generate a PDF of the physical connectivity graph.
 
-        i = 0
+        :param filename: The output PDF filename.
+        :param dif: The DIF to export.
+        """
         try:
-            for node in self.nodes:
-                g_node = pydot.Node(node.name,
-                                    label=node.name,
-                                    style="filled",
-                                    fillcolor=colors[i],
-                                    fontcolor=fcolors[i])
-
-                gvizg.add_node(g_node)
-                i += 1
-                if i == len(colors):
-                    i = 0
-        except:
-            logger.error('Failed to create pydot Node')
-            return
+            import pydot
 
-        try:
-            for dif in self.dif_ordering:
-                if isinstance(dif, topology.ShimEthDIF):
-                    edge = pydot.Edge(dif.members[0].name,
-                                      dif.members[1].name,
-                                      label=dif.name,
-                                      color="black")
-                    gvizg.add_edge(edge)
-        except:
-            logger.error('Failed to create pydot Edge')
-            return
+            colors = ['red', 'green', 'blue', 'orange', 'yellow']
+            fcolors = ['black', 'black', 'white', 'black', 'black']
+
+            gvizg = pydot.Dot(graph_type='digraph')
+
+            i = 0
+            try:
+                for node in dif.members:
+                    g_node = pydot.Node(node.name,
+                                        label=node.name,
+                                        style="filled",
+                                        fillcolor=colors[i],
+                                        fontcolor=fcolors[i])
+
+                    gvizg.add_node(g_node)
+                    i += 1
+                    if i == len(colors):
+                        i = 0
+            except Exception as e:
+                logger.error('Failed to create pydot Node: ' + str(e))
+                return
+
+            try:
+                for enroll, mgmt, dt in zip(self.enrollments,
+                                            self.mgmt_flows,
+                                            self.dt_flows):
+                    for e in enroll:
+                        if e['dif'] is not dif:
+                            continue
+
+                        edge = pydot.Edge(e['enrollee'].node.name,
+                                          e['enroller'].node.name,
+                                          color='black')
+                        gvizg.add_edge(edge)
+
+                    for e in mgmt:
+                        if e['src'].dif is not dif:
+                            continue
+
+                        edge = pydot.Edge(e['src'].node.name,
+                                          e['dst'].node.name,
+                                          color='blue')
+                        gvizg.add_edge(edge)
+
+                    for e in dt:
+                        if e['src'].dif is not dif:
+                            continue
+
+                        edge = pydot.Edge(e['src'].node.name,
+                                          e['dst'].node.name,
+                                          color='red')
+                        gvizg.add_edge(edge)
+            except Exception as e:
+                logger.error('Failed to create pydot Edge: ' + str(e))
+                return
+
+            try:
+                gvizg.write_pdf(filename)
+                logger.info('Generated PDF of DIF graph')
+            except Exception as e:
+                logger.error('Failed to write PDF: ' + str(e))
 
-        try:
-            gvizg.write_pdf(filename)
-            logger.info('Generated pdf of connectivity graph')
         except:
-            logger.error('Failed to write PDF')
+            logger.error('Warning: pydot module not installed, '
+                         'cannot produce DIF graph images')
 
 class Executor:
     __metaclass__ = abc.ABCMeta