Source code for schrodinger.test

# flake8: noqa
import contextlib
import inspect
import io
import logging
import os
import sys
import typing
from pathlib import Path

# For documentation, please see schrodinger/test/testfiles.h
from schrodinger.infra.test import internal_data_file
from schrodinger.infra.test import mmshare_data_file
from schrodinger.infra.test import mmshare_testfile
from schrodinger.utils.env import prepend_sys_path


def relpath_from_testdir(*path_components):
    """
    Absolute path to a file in the same directory as the calling code.
    """
    stk = inspect.stack()
    caller_path = stk[1][1]
    caller_dir = os.path.abspath(os.path.dirname(caller_path))
    return os.path.join(caller_dir, *path_components)


def assert_structures_identical(st1, st2, tolerance=1e-6):
    """
    Determines if the given two structures are exactly the same. Specifically,
    determines if the two structures have the same atoms with the same properties
    and bonds in the same order.

    :param st1: First structure to compare
    :type st1: Structure

    :param st2: Second structure to compare
    :type st2: Structure

    :param float tolerance: Absolute tolerance for comparing fractional atom and
                            bond properties. Relative tolerance is not supported
                            because properties read from file are limited to six
                            decimal places, so this assertion might return false
                            depending on the absolute value of the property.
    """
    # Prevent circular dependencies
    from schrodinger.test import custom_assertions

    custom_assertions.assert_properties_match(st1,
                                              st2,
                                              ignore=['i_m_ct_format'])
    custom_assertions.assertSameNumberOfAtoms(st1, st2)
    for at1, at2 in zip(list(st1.atom), list(st2.atom)):

        # Check at1 and at2 have the same properties
        custom_assertions.assertEqualDicts(dict(at1.property),
                                           dict(at2.property),
                                           tolerance=tolerance)

        # Check that at1 and at2 have the same bonds and bond types
        assert len(at1.bond) == len(at2.bond)
        prop_names = ('Atom 1 Index', 'Atom 2 Index', 'Bond Type')
        bonds1 = sorted([
            (int(b.atom1), int(b.atom2), b.type) for b in list(at1.bond)
        ])
        bonds2 = sorted([
            (int(b.atom1), int(b.atom2), b.type) for b in list(at2.bond)
        ])

        for b1, b2 in zip(bonds1, bonds2):
            for prop1, prop2, prop_name in zip(b1, b2, prop_names):
                assert prop1 == prop2, f"Bond {b1} in at1 and bond {b2} in at2 are different."

        a1_bonds = list(at1.bond)
        a2_bonds = list(at2.bond)

        a1_bonds.sort(key=lambda bnd: sorted([int(bnd.atom1), int(bnd.atom2)]))
        a2_bonds.sort(key=lambda bnd: sorted([int(bnd.atom1), int(bnd.atom2)]))

        for a1_bond, a2_bond in zip(a1_bonds, a2_bonds):
            custom_assertions.assertEqualDicts(a1_bond.property,
                                               a2_bond.property,
                                               tolerance=tolerance)


def mmshare_source_dir():
    return Path(os.environ["SCHRODINGER_SRC"]) / "mmshare"


def mmshare_nondist_dir():
    return mmshare_source_dir() / "python" / "nondist"


@contextlib.contextmanager
def add_build_tools_to_path():
    with prepend_sys_path(mmshare_source_dir() / "build_tools"):
        yield


@contextlib.contextmanager
def add_nondist_to_path():
    with prepend_sys_path(mmshare_nondist_dir()):
        yield


@contextlib.contextmanager
def stringio_handler(logger: logging.Logger) -> typing.Iterator[io.StringIO]:
    """
    Temporarily add a StringIO as the stream of a handler in order to assert
    output of log.
    """
    stream = io.StringIO()
    handler = logging.StreamHandler(stream)
    logger.addHandler(handler)
    yield stream
    logger.removeHandler(handler)