Source code for schrodinger.infra.mmerr

"""
Utility functions for the mmerr library.

Copyright Schrodinger, LLC. All rights reserved.

"""

from contextlib import ContextDecorator

from . import mm
# Import levels for convenience; accessible to user scripts:
from .mm import MMERR_DEBUG  # noqa: F401
from .mm import MMERR_FATAL  # noqa: F401
from .mm import MMERR_INFO  # noqa: F401
from .mm import MMERR_OFF
from .mm import MMERR_WARNING

_loglevel_stacks = {}


def push_level(mmerr_level, mmerr_handle=None):
    """
    Set the specified mmerr level and push the previous level onto a stack.

    :type mmerr_level: int
    :param mmerr_level: The logging level to use.

    :type mmerr_handle: int
    :param mmerr_handle: The mmerr handle to use. Defaults to
            mm.error_handler.

    """

    if mmerr_handle is None:
        mmerr_handle = mm.error_handler

    current_level = mm.mmerr_get_level(mmerr_handle)
    _loglevel_stacks.setdefault(mmerr_handle, []).append(current_level)
    mm.mmerr_level(mmerr_handle, mmerr_level)


def pop_level(mmerr_handle=None):
    """
    Pop the mmerr level stack and set logging level to the previous value.

    Return the mmerr level that has been set (i.e. the current logging level
    in use after the function returns).

    :type mmerr_handle: int
    :param mmerr_handle: The mmerr handle to use. Defaults to
            mm.error_handler.

    """

    if mmerr_handle is None:
        mmerr_handle = mm.error_handler

    previous_level = _loglevel_stacks[mmerr_handle].pop()
    mm.mmerr_level(mmerr_handle, previous_level)

    return previous_level


class Level(ContextDecorator):
    """
    Change mmerr level within a context or while a function is being run.
    Uses the default mm.error_handler level (MMERR_WARNING) if none is provided.
    Uses the default mm.error_handle if none is provided.
    """

    def __init__(self, level=MMERR_WARNING, handle=None):
        self.handle = handle
        self.level = level

    def __enter__(self):
        """Enter the context manager: changes error level for the specified error handler."""
        push_level(self.level, self.handle)

    def __exit__(self, exc_type, exc_value, traceback):
        """
        Exit context manager: return mmerr reporting to its previous level.

        Exception arguments are unused, but required by context manager
        signature.
        """
        pop_level(self.handle)


class disable_mmerr(Level):
    """
    Turn off mmerr reporting within a context or while a function is being run.
    Uses the default mm.error_handle if none is provided
    """

    def __init__(self, handle=None):
        super().__init__(level=MMERR_OFF, handle=handle)


class Capture(mm.CaptureMMErr):
    """
    Prefer using mm.CaptureMMErr directly.

    Capture all mmerr in a context. At the end of the context, messages will
    be available in the messages attribute. Messages are not printed.

    usage::

        with schrodinger.infra.CaptureMMErr() as captured:
            # make some mm calls
        print captured.messages

    """

    def __init__(self, handle=mm.error_handler):
        super().__init__(handle)