Source code for schrodinger.test.stu.outcomes.custom.phase_workups

"""
Special STU workups for testing Phase.

Copyright Schrodinger LLC, All Rights Reserved.
"""

import base64

from schrodinger.application.phase import constants
from schrodinger.infra import phase
from schrodinger.test.stu.outcomes import failures as failure_types
from schrodinger.test.stu.outcomes import standard_workups
from schrodinger.test.stu.outcomes import structure_comparisons
from schrodinger.utils.fileutils import tempfilename

BASE64_FAILURE_MSG = "%s: decoded base64 properties different"


def compare_phypo_files(check, reference, *properties, **tolerances):
    """
    Compare Phase hypothesis files. This is effectively the same calls as
    compare_mae_files, but with the .phypo extension. Phase hypothesis
    properties which are Base64-encoded strings are ignored in the comparison.
    """

    # Assure that both files have the .phypo extension
    msg = "File must have .phypo extension: %s"
    if not check.endswith(phase.PHASE_HYPO_FILE_EXT):
        raise failure_types.WorkupFailure(msg % check)
    if not reference.endswith(phase.PHASE_HYPO_FILE_EXT):
        raise failure_types.WorkupFailure(msg % reference)

    properties = structure_comparisons._move_args_to_kwargs(
        properties, tolerances)
    ignored_properties = [
        k for k, v in tolerances.items() if v == structure_comparisons.IGNORE
    ]

    # Ignore the hypothesis base64 encoded properties during mae comparison,
    # as well as phase_find_common project runtime path
    for propname in constants.base64_properties + [phase.PHASE_RUNNAME]:
        tolerances[propname] = structure_comparisons.IGNORE
    structure_comparisons.compare_mae_files(check, reference, *properties,
                                            **tolerances)

    # Compare base64 encoded properties of each Phase hypothesis file
    _compare_encoded_properties(check, reference, ignored_properties)

    # Return True if all comparisons successful (no exceptions raised)
    return True


def _compare_encoded_properties(check, reference, ignored_properties):
    """
    Compares base64 properties between two .phypo hypothesis files, by first
    decoding the properties to account for platform differences in econding.
    """

    failures = {}
    for i, st1, st2 in structure_comparisons._zip_structures(check, reference):
        # Check hypothesis base64 encoded properties
        base64_failures = []
        for prop in constants.base64_properties:
            if prop in ignored_properties:
                continue
            if not _equal_base64_property(prop, st1, st2):
                base64_failures.append(BASE64_FAILURE_MSG % prop)
        # Store any failures
        if base64_failures:
            failures["Structure %02d" % i] = base64_failures

    if failures:
        msg = ("Hypothesis base64 properties did not match in %s and %s." %
               (check, reference))
        raise structure_comparisons.CompositeFailure(msg, failures)


def _equal_base64_property(ct_property, st1, st2):
    """
    Returns whether the base-64 decoded property from two cts is equivalent. For
    float values embedded in any such property, values are compared to a given
    tolernace. If the property is not present in both cts, return False.

    :param ct_property: base-64 property name
    :type ct_property: str

    :param st1: structure ct
    :type st1: `structure.Structure`

    :param st2: structure ct
    :type st2: `structure.Structure`

    :return: whether the properties are equal
    :rtype: bool
    """

    value1 = st1.property.get(ct_property)
    value2 = st2.property.get(ct_property)

    # Ignore values not present in either ct
    if value1 is None and value2 is None:
        return True
    # Failure if one of the structures does not contain the given property
    if value1 is None or value2 is None:
        return False

    decoded_str1 = base64.b64decode(value1)
    decoded_str2 = base64.b64decode(value2)

    # Use standard workup to diff with default floating point tolerance
    with tempfilename() as file1, tempfilename() as file2:
        with open(file1, "wb") as fh:
            fh.write(decoded_str1)
        with open(file2, "wb") as fh:
            fh.write(decoded_str2)
        try:
            standard_workups.compare_files(file1, file2)
        except failure_types.WorkupFailure:
            return False
    return True