Add backward compatibility

src/moldrug/compat.py

@@ -0,0 +1,28 @@
+"""
+This module provides backward compatibility with moldrug <= 4.0.
+
+It allows pickle objects created with older versions of moldrug
+to be successfully loaded after internal refactoring (e.g. moved classes).
+"""
+import dill
+import importlib
+from typing import Dict, Tuple
+
+# Mapping from old (module, class) → new (module, class)
+_CLASS_RENAMES: Dict[Tuple[str, str], Tuple[str, str]] = {
+    ("moldrug.utils", "GA"): ("moldrug.opt", "GA"),
+    ("moldrug.utils", "Local"): ("moldrug.opt", "Local"),
+    # Add more mappings here as needed
+}
+
+
+class BackCompatUnpickler(dill.Unpickler):
+    """Custom unpickler that redirects renamed or moved classes."""
+
+    def find_class(self, module: str, name: str):
+        key = (module, name)
+        if key in _CLASS_RENAMES:
+            new_module, new_name = _CLASS_RENAMES[key]
+            mod = importlib.import_module(new_module)
+            return getattr(mod, new_name)
+        return super().find_class(module, name)

src/moldrug/opt.py

@@ -219,7 +219,7 @@ class GA:
     randomseed : Union[None, int]
         The random seed to use with random module.
     __moldrug_version__ : str
-        The molDrug version.
+        The moldrug version.
     costfunc : object
         The cost function set by the user.
     crem_db_path : str

src/moldrug/utils.py

@@ -19,7 +19,7 @@
                    RDKitMolCreate)
 from rdkit import Chem, RDLogger
 from rdkit.Chem import AllChem, DataStructs, Descriptors, Lipinski, rdFMCS
-
+from moldrug.compat import BackCompatUnpickler
 logger = logging.getLogger(__name__)
 
 RDLogger.DisableLog('rdApp.*')
@@ -465,6 +465,18 @@ def confgen(mol: Chem.rdchem.Mol, return_mol: bool = False, randomseed: Union[in
     else:
         return pdbqt_string
 
+    AllChem.EmbedMolecule(mol, randomSeed=randomSeed)
+    # The optimization introduce some sort of non-reproducible results.
+    # For that reason is not used when randomseed is set
+    if not randomseed:
+        AllChem.MMFFOptimizeMolecule(mol, maxIters=500)
+    preparator = MoleculePreparation()
+    mol_setups = preparator.prepare(mol)
+    pdbqt_string = PDBQTWriterLegacy.write_string(mol_setups[0])[0]
+    if return_mol:
+        return (pdbqt_string, mol)
+    else:
+        return pdbqt_string
 
 def update_reactant_zone(parent: Chem.rdchem.Mol, offspring: Chem.rdchem.Mol,
                          parent_replace_ids: List[int] = None, parent_protected_ids: List[int] = None):
@@ -841,9 +853,8 @@ def decompress_pickle(file: str):
     object
         The python object.
     """
-    data = bz2.BZ2File(file, 'rb')
-    data = pickle.load(data)
-    return data
+    with bz2.BZ2File(file, 'rb') as f:
+        return BackCompatUnpickler(f).load()
 
 
 def is_iter(obj):

streamlit/moldrug-dashboard.py

@@ -21,7 +21,7 @@
 
 import streamlit as st
 import streamlit.components.v1 as components
-from moldrug import utils
+from moldrug import utils, opt
 
 # TODO
 # add SyGma for metabolic prediction
@@ -333,10 +333,10 @@ def lig_prot_overview(_pop, protein_pdb_string):
 def load_pbz2(pbz2):
     moldrug_result = utils.decompress_pickle(pbz2)
     is_GA = False
-    if isinstance(moldrug_result, utils.GA):
+    if isinstance(moldrug_result, opt.GA):
         gen, pop = moldrug_result.NumGens, moldrug_result.pop
         is_GA = True
-    elif isinstance(moldrug_result, utils.Local):
+    elif isinstance(moldrug_result, opt.Local):
         gen, pop = 0, moldrug_result.pop
     elif isinstance(moldrug_result, tuple):
         if isinstance(moldrug_result[0], int) and isinstance(moldrug_result[1][0], utils.Individual):
@@ -608,7 +608,7 @@ def get_pubchem_dataframe(df: pd.DataFrame) -> pd.DataFrame:
                     if 'kept_gens' in grid.dataframe.columns:
                         props_to_drop.append('kept_gens')
 
-                    prop_df = grid.dataframe.drop(['mol', 'pdbqt', 'kept_gens', 'img', 'mols2grid-id'], axis=1).set_index('idx')
+                    prop_df = grid.dataframe.drop(props_to_drop, axis=1).set_index('idx')
                     st.download_button(
                         "Press to Download",
                         convert_df(prop_df),