Merge pull request #184 from bgyori/pysb_flat_annotations

Print model annotations and suppress extra newlines in pysb_flat export
pysb · Dec 8, 2015 · 15fac0c · 15fac0c
2 parents 21c8ae7 + 93fb45d
commit 15fac0c
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/pysb/export/pysb_flat.py b/pysb/export/pysb_flat.py
@@ -48,11 +48,12 @@ def export(self):
         output = StringIO()
 
         # Convenience function for writing out a componentset.
-        def write_cset(cset):
+        def write_cset(cset): 
             for c in cset:
                 output.write(repr(c))
                 output.write("\n")
-            output.write("\n")
+            if cset:
+                output.write("\n")
 
         if self.docstring:
             output.write('"""')
@@ -61,7 +62,8 @@ def write_cset(cset):
         output.write("# exported from PySB model '%s'\n" % self.model.name)
         output.write("\n")
         output.write("from pysb import Model, Monomer, Parameter, Expression, "
-                     "Compartment, Rule, Observable, Initial, ANY, WILD\n")
+                     "Compartment, Rule, Observable, Initial, Annotation, "
+                     "ANY, WILD\n")
         output.write("\n")
         output.write("Model()\n")
         output.write("\n")
@@ -75,5 +77,6 @@ def write_cset(cset):
         for pattern, value in self.model.initial_conditions:
             output.write("Initial(%s, %s)\n" % (repr(pattern), value.name))
         output.write("\n")
+        write_cset(self.model.annotations)
 
         return output.getvalue()