Python selection.parse_selection函数代码示例

def test_bool():
    sp = parse_selection("protein or water")
    eq(sp.source, "(atom.residue.is_protein or atom.residue.is_water)")

    sp = parse_selection("protein or water or all")
    eq(sp.source,
       "(atom.residue.is_protein or atom.residue.is_water or True)")

def test_nested_bool():
    sp = parse_selection("nothing and water or all")
    eq(sp.source,
       "((False and atom.residue.is_water) or True)")

    sp = parse_selection("nothing and (water or all)")
    eq(sp.source,
       "(False and (atom.residue.is_water or True))")

def test_unary_2():
    sp = parse_selection('all')
    for a in tt.atoms:
        assert sp.expr(a)

    sp = parse_selection('none')
    for a in tt.atoms:
        assert not sp.expr(a)

def test_not():
    sp = parse_selection("not protein")
    eq(sp.source, "(not atom.residue.is_protein)")

    sp = parse_selection("not not protein")
    eq(sp.source, "(not (not atom.residue.is_protein))")

    sp = parse_selection('!protein')
    eq(sp.source, "(not atom.residue.is_protein)")

def test_range():
    eq(parse_selection('resSeq 1 to 10').astnode, pnode('1 <= atom.residue.resSeq <= 10'))

    sp = parse_selection('resSeq 5 to 6')
    for a in tt.atoms:
        assert sp.expr(a)
    sp = parse_selection('resSeq 7 to 8')
    for a in tt.atoms:
        assert not sp.expr(a)

def test_values():
    sp = parse_selection("resid 4")
    eq(sp.source, "(atom.residue.index == 4)")

    sp = parse_selection("resid > 4")
    eq(sp.source, "(atom.residue.index > 4)")

    sp = parse_selection("resid gt 4")
    eq(sp.source, "(atom.residue.index > 4)")

    sp = parse_selection("resid 5 to 8")
    eq(sp.source, "(5 <= atom.residue.index <= 8)")

def test_binary_1():
    sp = parse_selection('resname "ALA"')
    assert sp.expr(tt.atom(0))
    assert sp.expr(tt.atom(1))

    sp = parse_selection('mass > 2')
    assert sp.expr(tt.atom(0))
    assert not sp.expr(tt.atom(1))
    assert sp.expr(tt.atom(2))

    sp = parse_selection('name ne O')
    assert sp.expr(tt.atom(0))
    assert not sp.expr(tt.atom(2))

def test_quotes():
    should_be = "((atom.name == 'CA') and (atom.residue.name == 'ALA'))"

    sp = parse_selection("name CA and resname ALA")
    eq(sp.source, should_be)
    assert sp.expr(tt.atom(0))

    sp = parse_selection('name "CA" and resname ALA')
    eq(sp.source, should_be)
    assert sp.expr(tt.atom(0))

    sp = parse_selection("name 'CA' and resname ALA")
    eq(sp.source, should_be)
    assert sp.expr(tt.atom(0))

def test_unary_1():
    eq(parse_selection('all').astnode, pnode('True'))
    eq(parse_selection('everything').astnode, pnode('True'))
    eq(parse_selection('none').astnode, pnode('False'))
    eq(parse_selection('nothing').astnode, pnode('False'))
    # eq(parse_selection('nucleic').astnode, pnode('atom.residue.is_nucleic'))
    # eq(parse_selection('is_nucleic').astnode, pnode('atom.residue.is_nucleic'))
    eq(parse_selection('protein').astnode, pnode('atom.residue.is_protein'))
    eq(parse_selection('is_protein').astnode, pnode('atom.residue.is_protein'))
    eq(parse_selection('water').astnode, pnode('atom.residue.is_water'))
    eq(parse_selection('is_water').astnode, pnode('atom.residue.is_water'))
    eq(parse_selection('waters').astnode, pnode('atom.residue.is_water'))

    def select(self, selection_string):
        """Execute a selection against the topology

        Parameters
        ----------
        selection_string : str
            An expression in the MDTraj atom selection DSL

        Examples
        --------
        >>> topology.select('name O and water')
        array([1, 3, 5, 10, ...])

        Returns
        -------
        indices : np.ndarray, dtype=int, ndim=1
            Array of the indices of the atoms matching the selection expression.

        See Also
        --------
        select_expression, mdtraj.core.selection.parse_selection
        """

        filter_func = parse_selection(selection_string).expr
        indices = np.array([a.index for a in self.atoms if filter_func(a)])
        return indices

def test_unary_3():
    sp = parse_selection('protein or water')

    for a in tt.atoms:
        assert sp.expr(a)

    sp = parse_selection('protein and water')
    for a in tt.atoms:
        assert not sp.expr(a)

    sp = parse_selection('not (protein and water)')
    for a in tt.atoms:
        assert sp.expr(a)

    sp = parse_selection('not not (protein and water)')
    for a in tt.atoms:
        assert not sp.expr(a)

def test_raises2():
    pytest.raises(ValueError, lambda: parse_selection('dog 5'))
    pytest.raises(ValueError, lambda: parse_selection('dog == 5'))
    pytest.raises(ValueError, lambda: parse_selection('dog frog'))
    pytest.raises(ValueError, lambda: parse_selection('not dog'))
    pytest.raises(ValueError, lambda: parse_selection('protein or dog'))
    pytest.raises(ValueError, lambda: parse_selection('dog 1 to 5'))
    pytest.raises(ValueError, lambda: parse_selection('dog'))

    def select_expression(self, selection_string):
        """Translate a atom selection expression into a pure python expression.

        Parameters
        ----------
        selection_string : str
            An expression in the MDTraj atom selection DSL

        Examples
        --------
        >>> topology.select_expression('name O and water')
        "[atom.index for atom in topology.atoms if ((atom.name == 'O') and atom.residue.is_water)]")

        Returns
        -------
        python_string : str
            A string containing a pure python expression, equivalent to the
            selection expression.
        """
        condition = parse_selection(selection_string).source
        fmt_string = "[atom.index for atom in topology.atoms if {condition}]"
        return fmt_string.format(condition=condition)

def test_raises():
    pytest.raises(ValueError, lambda: parse_selection('or'))
    pytest.raises(ValueError, lambda: parse_selection('a <'))

def test_in():
    sp = parse_selection("resname ALA ASP GLU")
    eq(sp.source, "(atom.residue.name in ['ALA', 'ASP', 'GLU'])")

    sp = parse_selection("resid 100 101 102")
    eq(sp.source, "(atom.residue.index in [100, 101, 102])")