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

import os

from hypothesis import strategies

from schrodinger import structure
from schrodinger.test import mmshare_testfile

BLACKLISTED_SUBDIRS = ('error_examples',)


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):
        dirnames[:] = [
            dirname for dirname in dirnames
            if dirname not in BLACKLISTED_SUBDIRS
        ]
        for fname in filenames:
            if fname.endswith(extensions):
                structure_files.append(os.path.join(dirpath, fname))

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


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


[docs]@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))
[docs]@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))