Source code for schrodinger.application.steps.licensing

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 `LICENSE_KEY`.

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

    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