Source code for schrodinger.application.desmond.picklejar

"""
Facilities for pickling objects that are otherwise uneasy to pickle, e.g.,
class attributes.

Copyright Schrodinger, LLC. All rights reserved.

"""

# Contributors: Yujie Wu

import pickle

_BASE_PATH = 'schrodinger.application.desmond.'
_STAGE_PATH = _BASE_PATH + 'stage'

_relative_unpickle_map = {
    'restraint.restraint_stage': 'stage.prepare.forcefield',
    'fragment_linking.stage': 'stage.app.fragment_linking.stage',
    'mxmd.mxmd_stage': 'stage.app.mxmd',
    'absolute_binding.stage': 'stage.app.absolute_binding.stage'
}

_absolute_unpickle_map = {
    _BASE_PATH + k: _BASE_PATH + v for (k, v) in _relative_unpickle_map.items()
}


# We need a user defined object in order to be able to set attributes when
# unpickling
class _UserDefinedObject:

    def __setstate__(self, _):
        # Fix "UnpicklingError: state is not a dictionary"
        pass


_unpickle_deprecated_names = {
    'Trajectory', 'GCMC', 'SolvatePocketBase', 'Minimize', 'McproStageBase',
    'McproSystemBuilder', 'McproSimulate', 'FepAnalysis2', 'WaterMap',
    'RetryTimer', 'BindingPoseFepPrimer'
}

_unpickle_deprecated_modules = {
    'schrodinger.application.desmond.stage.app.binding_pose'
}


class CustomUnpickler(pickle.Unpickler):

    def find_class(self, module, name):
        remapped_module = _absolute_unpickle_map.get(module, module)
        try:
            return super().find_class(remapped_module, name)
        except AttributeError:
            # check this after trying super in case someone defines a new
            # class with an old name
            if name in _unpickle_deprecated_names:
                return _UserDefinedObject
            raise
        except ModuleNotFoundError:
            if remapped_module in _unpickle_deprecated_modules:
                return _UserDefinedObject
            raise


class PickleJar(object):
    """

    """

    jar = {}

    @staticmethod
    def _update_data():
        """

        """
        jar = PickleJar.jar
        for k in jar:
            for dn in jar[k]:
                jar[k][dn] = k.__dict__[dn]

    @staticmethod
    def serialize(fh):
        """

        """
        PickleJar._update_data()
        pickle.dump(PickleJar.jar, fh)

    @staticmethod
    def deserialize(fh):
        """

        """
        unpickler = CustomUnpickler(fh)
        jar = unpickler.load()
        for k in jar:
            for dn in jar[k]:
                setattr(k, dn, jar[k][dn])


class PicklableMetaClass(type):

    def __init__(cls, name, bases, dict):
        PickleJar.jar[cls] = {}
        jar = PickleJar.jar[cls]
        for k in dict:
            v = dict[k]
            if (isinstance(v, tuple)):
                try:
                    if (v[1] == Picklable):
                        setattr(cls, k, v[0])
                        jar[k] = None
                except IndexError:
                    pass


class Picklable(object, metaclass=PicklableMetaClass):
    pass


class PickleState(object):
    pass