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

import sys

from schrodinger import structure
from schrodinger.structutils import analyze


def asl_count(asl, expression, fn, ct_index=-1):
    """
Return False if the # of atoms doesn't match the expression given

:param asl: asl expression
:type  asl: string
:param expression: expression for how the count should compare (e.g. '>=5')
:type  expression: string
:param fn: maestro filename
:type  fn: string
:param ct_index: Specific index of ct to be tested, otherwise all will be
:type  ct_index: str, int, or float (converted to int)
    Example:
    custom.asl_count.asl_count("water and atom.ele O", "==0", "mol.mae")

"""
    # create test function, note 'eval' is not safe
    # >=, <=, >, <, ==, !=
    if ">=" in expression:
        test = lambda x: x >= float(expression.split('>=')[1])
    elif "<=" in expression:
        test = lambda x: x <= float(expression.split('<=')[1])
    elif ">" in expression:  # this must come after >= in code!
        test = lambda x: x > float(expression.split('>')[1])
    elif "<" in expression:
        test = lambda x: x < float(expression.split('<')[1])
    elif '==' in expression:
        test = lambda x: x == float(expression.split('==')[1])
    elif '!=' in expression:
        test = lambda x: x != float(expression.split('!=')[1])
    else:
        raise NotImplementedError(
            f"Error, this expression is not implemented: {expression}")
    ct_index = int(ct_index)
    errors = []
    for ict, ct in enumerate(structure.StructureReader(fn)):
        if ct_index in [-1, ict]:
            count = len(analyze.evaluate_asl(ct, asl))
            passed = test(count)
            if not passed:
                msg = "# of atoms ({}) in ct index {} matching '{}' is not '{}'.".format(
                    count, ict, asl, expression)
                errors.append(msg)
    if errors:
        raise AssertionError('\n'.join(errors))
    return True


if __name__ == "__main__":
    """
    Usage
    $SCHRODINGER/run asl_count.py <asl> <filename>
    """
    print(sys.argv)
    if (asl_count(*sys.argv[1:])):
        print("PASS")
    else:
        print("FAIL")