Source code for schrodinger.application.steps.licensing

from schrodinger.tasks import stepper
from schrodinger.utils import license


class LicenseMixin:
    """
    A mixin to ensure that a step uses a `LICENSE_KEY` license.

    To be mixed in with a `stepper.ReduceStep`.
    Subclasses need to define `NO_LICENSE_MSG` and `LICENSE_KEY`.

    If no licenses are available when the step is executed an `EXCEPTION` will
    be raised.
    """
    LICENSE_KEY = None
    NO_LICENSE_MSG = None
    NO_MORE_LICENSES = 'No license tokens available.'
    EXCEPTION = RuntimeError

    def _hasLicense(self):
        lic = license.License(self.LICENSE_KEY)
        try:
            valid = lic.isValid()
        finally:
            lic.checkin()
        return valid

    def validateSettings(self):
        """
        In addition to the regular settings validation check for the license
        validity.

        :return: a list of validation issues
        :rtype: List[stepper.ValidationIssue]
        """
        issues = []
        if not self._hasLicense():
            issues = [stepper.SettingsError(self, self.NO_LICENSE_MSG)]
        return issues + super().validateSettings()

    def reduceFunction(self, inputs):
        """
        Perform the step with checking out (and back in) of the license.

        See also `stepper.ReduceStep.reduceFunction`.
        """
        lic = license.License(self.LICENSE_KEY)
        try:
            if not lic.isValid():
                raise self.EXCEPTION(self.NO_MORE_LICENSES)
            outputs = super().reduceFunction(inputs)
        finally:
            lic.checkin()
        return outputs

    def getLicenseRequirements(self):
        required_licenses = super().getLicenseRequirements()
        required_licenses[self.LICENSE_KEY] = 1
        return required_licenses