diff --git a/pyteomics/proforma.py b/pyteomics/proforma.py index c24792b..3a17f58 100644 --- a/pyteomics/proforma.py +++ b/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 diff --git a/tests/test_proforma.py b/tests/test_proforma.py index 6043af8..2425f83 100644 --- a/tests/test_proforma.py +++ b/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()