Source code for schrodinger.test.stu.workups.workup_tools

import functools

messages = []


def workup(workup_function):
    """
    Wrap a STU workup to:
        * Return False for AssertionError
        * Raise a WorkupImportFailure if the workup function does not exist
        * Return True for return None?
        * Capture messages and log to 'stu.workup'

    """

    @functools.wraps(workup_function)
    def wrapped(*args, **kwargs):
        err = None
        try:
            rvalue = workup_function(*args, **kwargs)
            if rvalue is None:
                rvalue = True
        except AssertionError as e:
            err = e  # py3 needs reassignment for scope
            rvalue = False
        if rvalue:
            return rvalue

        display_args = _display_args(args, kwargs)

        message = workup_function.__name__
        message += "("
        message += ', '.join(display_args)
        message += ") failed"
        if err:
            message += f":\n{err}"
        messages.append(message)
        return rvalue

    setattr(wrapped, "wrapped", True)
    return wrapped


def _display_args(args, kwargs):
    display_args = list(map(repr, args))
    # uses sorted to make tests consistent
    display_args.extend(f'{k}={v!r}' for k, v in sorted(kwargs.items()))
    return display_args