Source code for schrodinger.models.advanced_mappers

from schrodinger.models import mappers


class ModelBaseClassMapperMixin(mappers.MapperMixin):
    """
    A mapper mixin that allows any model class that is a subclass of a
    specified base class. Unlike a standard MapperMixin, the mapper is re-
    generated every time a new model is set, and the model_class class variable
    is updated to the class of the new model.

    This allows mappings to vary depending on the characteristics of the model
    class (for example, values of class attributes, presence or absence of
    specific params, etc.).

    To use, set the model_base_class class variable to the base class of all
    models that will be used with this mapper.
    """
    model_base_class = None

    def _setupMapperMixin(self):
        if self.model_class is not None:
            raise ValueError('The model_class class variable should not be '
                             'defined for a ModelBaseClassMapperMixin %s. '
                             'Define the model_classes tuple instead.' %
                             self.__class__.__name__)
        self.model_class = self.model_base_class
        super()._setupMapperMixin()

    def _setupModelClass(self, model):
        super()._setupModelClass(model)
        if not isinstance(model, self.model_base_class):
            raise TypeError(f'Model must be a subclass of '
                            f'{self.model_base_class}. Got {type(model)}')
        self.model_class = type(model)
        self._buildMapper()


class MultiModelClassMapperMixin(mappers.MapperMixin):
    """
    A mapper mixin that allows multiple model classes. To use, define the
    model_classes tuple with all the allowable model classes. Then define the
    usual defineMappings and getSignalsAndSlots methods, which are allowed to
    branch on self.model_class to specify model class-specific mappings and
    signal/slot connections.

    In use, self.model_class will be set to one of the classes in
    self.model_classes, depending on what the current model is.
    """

    model_classes = tuple()

    def _setupMapperMixin(self):
        if self.model_class is not None:
            raise ValueError('The model_class class variable should not be '
                             'defined for a MultiModelClassMapperMixin %s. '
                             'Define the model_classes tuple instead.' %
                             self.__class__.__name__)
        if not self.model_classes:
            raise ValueError('The model_classes variable must be non-empty.')
        self.model_class = self.model_classes[0]
        super()._setupMapperMixin()

    def _setupModelClass(self, model):
        super()._setupModelClass(model)
        for model_class in self.model_classes:
            if isinstance(model, model_class):
                self.model_class = model_class
                break
        else:
            raise TypeError(f'Model must be one of {self.model_classes}. '
                            f'Got {type(model)}')
        self._buildMapper()