Source code for schrodinger.infra.mminit

"""
Assist in initialization and termination of mmlibs.

"""
# Copyright Schrodinger, LLC. All rights reserved.

import contextlib
import functools

from schrodinger.infra import mm


class Initializer:
    """
    A class to handle initialization and termination of needed mmlibs.

    It is designed to be created on import and initialize needed mmlibs for
    a module. In general, automatic termination is not supported, but terminate
    method can be called manually.

    """

    def __init__(self, initializers, terminators):
        self.initializers = initializers
        self.terminators = terminators
        self.error_handler = None
        self.init_count = 0
        self.initialize()

    def initialize(self, error_handler=None):
        """
        Initialize all needed mmlibs.

        """
        if error_handler is None:
            error_handler = mm.error_handler

        # Keep track of the error handler used to initialize.
        self.error_handler = error_handler
        for init in self.initializers:
            init(error_handler)
        self.init_count += 1

    def terminate(self):
        """
        Terminate all previously initialized mmlibs.

        """
        for term in self.terminators:
            term()
        self.init_count -= 1

    def __copy__(self):
        # This implementation just returns 'self', effectively preventing
        # copying. Each Initializer instance should be a singleton providing
        # for initialization and termination of its set of mmlibs.
        #
        # I can't think of a reason to need an actual copy, but if a need
        # arises, be sure that init_count in the new instance is reset to
        # zero.
        return self

    def __deepcopy__(self, memo=None):
        # Prevent copying. See discussion in __copy__.
        return self


class _mmlib_new_delete:
    """
    Use new and delete as context manager pairs.

    Useful if new returns an `int`, which of course has no
    knowledge of the underlying object's lifetime.
    """

    def __init__(self, new, delete):
        self._new = new
        self._delete = delete

    def new(self, *args):
        v = self._new(*args)

        @contextlib.contextmanager
        def scope():
            try:
                yield v
            finally:
                self._delete(v)

        return scope()


@contextlib.contextmanager
def MMLib(init, term, new=None, delete=None):
    """
    Decorator and context manager to control initialization of a single
    library.

    """
    init(mm.error_handler)
    try:
        obj = None
        if new and delete:
            obj = _mmlib_new_delete(new, delete)
        elif new or delete:
            raise ValueError(
                'Please provide both `new` and `delete`, or neither.')
        yield obj
    finally:
        term()


# Make life easier by canning definitions for specific mmlibs:
#
# Use like::
#
#    with mminit.mmfrag() as mmfrag:
#       with mmfrag.new() as h:
#           # do stuff with `h`
#
mmct = functools.partial(MMLib, mm.mmct_initialize, mm.mmct_terminate)
mmfrag = functools.partial(MMLib, mm.mmfrag_initialize, mm.mmfrag_terminate,
                           mm.mmfrag_new, mm.mmfrag_delete)