Merge pull request #216 from jonls/write-compartments

native: Allow writing compartments in ModelWriter

psamm/datasource/native.py

@@ -1072,6 +1072,28 @@ def _dump(self, stream, data):
         finally:
             dumper.dispose()
 
+    def convert_compartment_entry(self, compartment, adjacencies):
+        """Convert compartment entry to YAML dict.
+
+        Args:
+            compartment: :class:`psamm.datasource.entry.CompartmentEntry`.
+            adjacencies: Sequence of IDs or a single ID of adjacent
+                compartments (or None).
+        """
+        d = OrderedDict()
+        d['id'] = compartment.id
+        if adjacencies is not None:
+            d['adjacent_to'] = adjacencies
+
+        order = {key: i for i, key in enumerate(['name'])}
+        prop_keys = set(compartment.properties)
+        for prop in sorted(prop_keys,
+                           key=lambda x: (order.get(x, 1000), x)):
+            if compartment.properties[prop] is not None:
+                d[prop] = compartment.properties[prop]
+
+        return d
+
     def convert_compound_entry(self, compound):
         """Convert compound entry to YAML dict."""
         d = OrderedDict()
@@ -1108,25 +1130,27 @@ def is_equation_valid(equation):
         prop_keys = (set(reaction.properties) -
                      {'lower_flux', 'upper_flux', 'reversible'})
         for prop in sorted(prop_keys, key=lambda x: (order.get(x, 1000), x)):
-            if reaction.properties[prop] is not None:
-                d[prop] = reaction.properties[prop]
+            if reaction.properties[prop] is None:
+                continue
+            d[prop] = reaction.properties[prop]
             if prop == 'equation' and not is_equation_valid(d[prop]):
                 del d[prop]
 
         return d
 
-    def write_compounds(self, stream, compounds, properties=None):
-        """Write iterable of compounds as YAML object to stream.
+    def _write_entries(self, stream, entries, converter, properties=None):
+        """Write iterable of entries as YAML object to stream.
 
         Args:
             stream: File-like object.
-            compounds: Iterable of compound entries.
-            properties: Set of compound properties to output (or None to output
-                all).
+            entries: Iterable of entries.
+            converter: Conversion function from entry to YAML object.
+            properties: Set of compartment properties to output (or None to
+                output all).
         """
         def iter_entries():
-            for c in compounds:
-                entry = self.convert_compound_entry(c)
+            for c in entries:
+                entry = converter(c)
                 if entry is None:
                     continue
                 if properties is not None:
@@ -1137,6 +1161,35 @@ def iter_entries():
 
         self._dump(stream, list(iter_entries()))
 
+    def write_compartments(self, stream, compartments, adjacencies,
+                           properties=None):
+        """Write iterable of compartments as YAML object to stream.
+
+        Args:
+            stream: File-like object.
+            compartments: Iterable of compartment entries.
+            adjacencies: Dictionary mapping IDs to adjacent compartment IDs.
+            properties: Set of compartment properties to output (or None to
+                output all).
+        """
+        def convert(entry):
+            return self.convert_compartment_entry(
+                entry, adjacencies.get(entry.id))
+
+        self._write_entries(stream, compartments, convert, properties)
+
+    def write_compounds(self, stream, compounds, properties=None):
+        """Write iterable of compounds as YAML object to stream.
+
+        Args:
+            stream: File-like object.
+            compounds: Iterable of compound entries.
+            properties: Set of compound properties to output (or None to output
+                all).
+        """
+        self._write_entries(
+            stream, compounds, self.convert_compound_entry, properties)
+
     def write_reactions(self, stream, reactions, properties=None):
         """Write iterable of reactions as YAML object to stream.
 
@@ -1146,15 +1199,5 @@ def write_reactions(self, stream, reactions, properties=None):
             properties: Set of reaction properties to output (or None to output
                 all).
         """
-        def iter_entries():
-            for r in reactions:
-                entry = self.convert_reaction_entry(r)
-                if entry is None:
-                    continue
-                if properties is not None:
-                    entry = OrderedDict(
-                        (key, value) for key, value in iteritems(entry)
-                        if key == 'id' or key in properties)
-                yield entry
-
-        self._dump(stream, list(iter_entries()))
+        self._write_entries(
+            stream, reactions, self.convert_reaction_entry, properties)

psamm/tests/test_datasource_native.py

@@ -14,7 +14,7 @@
 # You should have received a copy of the GNU General Public License
 # along with PSAMM.  If not, see <http://www.gnu.org/licenses/>.
 #
-# Copyright 2014-2015  Jon Lund Steffensen <jon_steffensen@uri.edu>
+# Copyright 2014-2017  Jon Lund Steffensen <jon_steffensen@uri.edu>
 
 import os
 import shutil
@@ -564,6 +564,19 @@ class TestNativeModelWriter(unittest.TestCase):
     def setUp(self):
         self.writer = native.ModelWriter()
 
+    def test_convert_compartment_entry(self):
+        compartment = entry.DictCompartmentEntry({
+            'id': 'cytosol',
+            'name': 'Cytosol',
+            'custom': 456
+        })
+        d = self.writer.convert_compartment_entry(compartment, 'periplasm')
+        self.assertIsInstance(d, OrderedDict)
+        self.assertEqual(d['id'], 'cytosol')
+        self.assertEqual(d['name'], 'Cytosol')
+        self.assertEqual(d['custom'], 456)
+        self.assertEqual(d['adjacent_to'], 'periplasm')
+
     def test_convert_compound_entry(self):
         compound = entry.DictCompoundEntry({
             'id': 'c1',
@@ -598,6 +611,41 @@ def test_convert_reaction_entry(self):
         self.assertEqual(d['custom_property'], -34)
         self.assertNotIn('another_custom_property', d)
 
+    def test_write_compartments(self):
+        stream = StringIO()
+        compartments = [
+            entry.DictCompartmentEntry({
+                'id': 'cytosol',
+                'name': 'Cytosol',
+                'custom': 456
+            }),
+            entry.DictCompartmentEntry({
+                'id': 'periplasm',
+                'name': 'Periplasm',
+                'test': 'abc'
+            })
+        ]
+        adjacencies = {
+            'cytosol': 'periplasm',
+            'periplasm': ['cytosol', 'e']
+        }
+        self.writer.write_compartments(stream, compartments, adjacencies)
+
+        self.assertEqual(yaml.safe_load(stream.getvalue()), [
+            {
+                'id': 'cytosol',
+                'adjacent_to': 'periplasm',
+                'name': 'Cytosol',
+                'custom': 456
+            },
+            {
+                'id': 'periplasm',
+                'adjacent_to': ['cytosol', 'e'],
+                'name': 'Periplasm',
+                'test': 'abc'
+            }
+        ])
+
     def test_write_compounds(self):
         stream = StringIO()
         compounds = [