Globally set RDKit to include all properties when pickling an RDKit Mol

gufe/components/explicitmoleculecomponent.py

@@ -125,6 +125,16 @@ def __init__(self, rdkit: RDKitMol, name: str = ""):
         self._smiles: Optional[str] = None
         self._name = name
 
+    def __getstate__(self):
+        # TODO: check that RDKit setting is set before issuing warning
+        if Chem.GetDefaultPickleProperties() != Chem.PropertyPickleOptions.AllProps:
+            warnings.warn(
+                "RDKit does not preserve Mol properties when pickled by default, which may drop e.g. atom charges; "
+                "consider setting `Chem.SetDefaultPickleProperties(Chem.PropertyPickleOptions.AllProps)`"
+            )
+
+        return self.__dict__
+
     @classmethod
     def _defaults(cls):
         return super()._defaults()

gufe/tests/test_explicitmoleculecomponent.py

@@ -0,0 +1,15 @@
+import pickle
+
+
+class ExplicitMoleculeComponentMixin:
+
+    def test_pickle(self, instance):
+
+        pickled = pickle.dumps(instance)
+        unpickled = pickle.loads(pickled)
+
+        assert unpickled == instance
+
+        # it's currently the case that the flyweight pattern isn't respected by
+        # pickling/unpickling
+        assert not unpickled is instance

gufe/tests/test_proteincomponent.py

@@ -13,6 +13,7 @@
 from gufe import ProteinComponent
 
 from .conftest import ALL_PDB_LOADERS
+from .test_explicitmoleculecomponent import ExplicitMoleculeComponentMixin
 from .test_tokenization import GufeTokenizableTestsMixin
 
 
@@ -81,7 +82,7 @@ def assert_topology_equal(ref_top, top):
         assert ref_bond[1].index == bond[1].index
 
 
-class TestProteinComponent(GufeTokenizableTestsMixin):
+class TestProteinComponent(GufeTokenizableTestsMixin, ExplicitMoleculeComponentMixin):
 
     cls = ProteinComponent
     repr = "ProteinComponent(name=Steve)"

gufe/tests/test_smallmoleculecomponent.py

@@ -24,6 +24,7 @@
 from gufe.components.explicitmoleculecomponent import _ensure_ofe_name
 from gufe.tokenization import TOKENIZABLE_REGISTRY
 
+from .test_explicitmoleculecomponent import ExplicitMoleculeComponentMixin
 from .test_tokenization import GufeTokenizableTestsMixin
 
 
@@ -74,7 +75,7 @@ def test_ensure_ofe_name(internal, rdkit_name, name, expected, recwarn):
     assert rdkit.GetProp("ofe-name") == out_name
 
 
-class TestSmallMoleculeComponent(GufeTokenizableTestsMixin):
+class TestSmallMoleculeComponent(GufeTokenizableTestsMixin, ExplicitMoleculeComponentMixin):
 
     cls = SmallMoleculeComponent
     repr = "SmallMoleculeComponent(name=ethane)"

news/pickle-explicitmoleculecomponent.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* Added warning when pickling an ``ExplicitMoleculeComponent`` that RDKit mol properties not preserved by default.
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>