Source code for schrodinger.test.hypothesis.strategies.structures

import os
import pathlib

from hypothesis import strategies

from schrodinger import structure
from schrodinger.test import mmshare_testfile

BLACKLISTED_SUBDIRS = (
    # these files are broken, we use them in tests that check that the proper
    # exceptions and error messages are issued in such cases
    'error_examples',
    # same as the previous one, but in the LD preprocessor
    'livedesign/preprocessor/error_examples',
    # these sdf files include queries, which mmmdl does not support
    'livedesign/substructure')


def _normalize_path(filepath):
    return str(pathlib.Path(filepath))


def _is_blacklisted_path(dirpath):
    return any(True for subpath in BLACKLISTED_SUBDIRS
               if _normalize_path(dirpath) == _normalize_path(
                   mmshare_testfile(subpath)))


def _get_structure_files():
    """
    :return: A list of the structure containing files in the testfiles
             directory, excluding those in blacklisted subdirectories.
    :rtype: list
    """
    structure_files = []
    testfiles = mmshare_testfile('')
    extensions = ('.mae', '.maegz', '.mae.gz', '.pdb', '.pdb.gz', '.sdf')
    for dirpath, dirnames, filenames in os.walk(testfiles, topdown=True):
        if _is_blacklisted_path(dirpath):
            dirnames[:] = []
            continue
        for fname in filenames:
            if fname.endswith(extensions):
                structure_files.append(os.path.join(dirpath, fname))

    assert len(structure_files) > 1000  # sanity check that we're really
    # collecting files here
    return structure_files


_STRUCTURE_FILES = _get_structure_files()  # So that we only do this once


@strategies.composite
def structure_files(draw):
    """
    :return: The full path of a structure-containing file in the testfiles
             directory
    :rtype: str
    """
    return draw(strategies.sampled_from(_STRUCTURE_FILES))


@strategies.composite
def structures(draw):
    """
    :return: A structure drawn from one of the files in the testfiles directory
    :rtype: structure.Structure
    """
    filename = draw(structure_files())
    cts = [ct for ct in structure.StructureReader(filename)]
    return draw(strategies.sampled_from(cts))