elements: Allow exporting to DMS YAML

This allows for exporting the experiment to a YAML format which is
parsable by the DMS.

2 files changed, 27 insertions, 0 deletions

diff --git a/rumba/elements/experimentation.py b/rumba/elements/experimentation.py
index 9890975..2c6f478 100644
--- a/rumba/elements/experimentation.py
+++ b/rumba/elements/experimentation.py
@@ -521,6 +521,15 @@ class Experiment(object):
         end = time.time()
         logger.info("Swap-out took %.2f seconds", end - start)
 
+    def to_dms_yaml(self, filename, clean=True):
+        mode = 'w'
+        if not clean:
+            mode += '+'
+        with open(filename, mode) as f:
+            for node in self.nodes:
+                f.write("---\n")
+                node.to_dms_yaml(f)
+            f.write("...\n")
 
 class Executor:
     __metaclass__ = abc.ABCMeta
diff --git a/rumba/elements/topology.py b/rumba/elements/topology.py
index c47d541..46f24c2 100644
--- a/rumba/elements/topology.py
+++ b/rumba/elements/topology.py
@@ -654,6 +654,24 @@ class Node(object):
     def __neq__(self, other):
         return not self == other
 
+    def to_dms_yaml(self, buffer):
+        buffer.write("node: %s\n" % (self.name))
+        buffer.write("registrations:\n")
+
+        for dif in self.dif_registrations:
+            if isinstance(dif, NormalDIF):
+                buffer.write(" - %s:\n" % (dif.name))
+                for reg in self.dif_registrations[dif]:
+                    tp = "normal"
+                    nm = reg.name
+                    if isinstance(reg, ShimEthDIF):
+                        tp = "eth"
+                        for member in reg.members:
+                            if member.name is not self.name:
+                                nm = member.name
+                    buffer.write("  - %sdifid: %s, diftype: %s, "
+                                 "diftype_number: 1%s\n" % ("{", nm, tp, "}"))
+
     def add_dif(self, dif):
         """
         Adds a DIF to the list.