From f9f749d4690feaf076542363fc4d9fdd98ddd61a Mon Sep 17 00:00:00 2001
From: Joshua Klein <mobiusklein@gmail.com>
Date: Thu, 28 Mar 2024 09:11:00 -0400
Subject: [PATCH 1/2] Adjust pickling

---
 pyteomics/proforma.py  | 15 +++++++++++++--
 tests/test_proforma.py | 13 +++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/pyteomics/proforma.py b/pyteomics/proforma.py
index c24792ba..80497249 100644
--- a/pyteomics/proforma.py
+++ b/pyteomics/proforma.py
@@ -654,6 +654,17 @@ 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):
+        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 +995,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 6043af82..2425f83a 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()

From 9fdb44f42ae323c9fe76a2fc54027eb4055eb8cb Mon Sep 17 00:00:00 2001
From: Joshua Klein <mobiusklein@gmail.com>
Date: Thu, 28 Mar 2024 09:16:16 -0400
Subject: [PATCH 2/2] Incomplete state

---
 pyteomics/proforma.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pyteomics/proforma.py b/pyteomics/proforma.py
index 80497249..3a17f586 100644
--- a/pyteomics/proforma.py
+++ b/pyteomics/proforma.py
@@ -658,6 +658,8 @@ 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