Merge pull request #144 from mobiusklein/fix/proforma_mod_pickle

Adjust pickling

pyteomics/proforma.py

@@ -654,6 +654,19 @@ def __init__(self, value, extra=None, group_id=None, style=None):
         self._definition = None
         self.style = style
 
+    def __reduce__(self):
+        return self.__class__, (self.value, self.extra, self.group_id, self.style), self.__getstate__()
+
+    def __getstate__(self):
+        if self._definition is None:
+            return None
+        state = self._definition.copy()
+        state['source'] = None
+        return state
+
+    def __setstate__(self, state):
+        self._definition = state
+
     def __eq__(self, other):
         if isinstance(other, ModificationToken):
             return other == self
@@ -984,9 +997,9 @@ class GenericModification(ModificationBase):
         partial(UnimodModification.resolver, strict=False)
     ])
 
-    def __init__(self, value, extra=None, group_id=None):
+    def __init__(self, value, extra=None, group_id=None, style=None):
         super(GenericModification, self).__init__(
-            value, extra, group_id)
+            value, extra, group_id, style)
 
     def _format_main(self):
         return self.value

tests/test_proforma.py

@@ -1,5 +1,6 @@
 from os import path
 import unittest
+import pickle
 import pyteomics
 pyteomics.__path__ = [path.abspath(
     path.join(path.dirname(__file__), path.pardir, 'pyteomics'))]
@@ -180,5 +181,17 @@ def test_mass_modification(self):
         self.assertIn(mod2, container)
 
 
+class ModificationPicklingTest(unittest.TestCase):
+    def test_pickle(self):
+        mod = GenericModification("UNIMOD:1")
+        payload = pickle.dumps(mod)
+        dup = pickle.loads(payload)
+        self.assertEqual(mod, dup)
+        assert mod.mass is not None
+        payload = pickle.dumps(mod)
+        dup = pickle.loads(payload)
+        self.assertEqual(mod, dup)
+
+
 if __name__ == '__main__':
     unittest.main()