def test_filter_with_predicate(self): x = Keyword(name='x') predicate = lambda item: item.name == 'x' body = Body(items=[Keyword(), x, Keyword()]) assert_equal(body.filter(predicate=predicate), [x]) body = Body(items=[Keyword(), If(), x, For(), Keyword()]) assert_equal(body.filter(keywords=True, predicate=predicate), [x])
def test_filter(self): k1, k2, k3 = Keyword(), Keyword(), Keyword() f1, i1, i2 = For(), If(), If() body = Body(items=[k1, f1, i1, i2, k2, k3]) assert_equal(body.filter(keywords=True), [k1, k2, k3]) assert_equal(body.filter(keywords=False), [f1, i1, i2]) assert_equal(body.filter(ifs=True, fors=True), [f1, i1, i2]) assert_equal(body.filter(ifs=False, fors=False), [k1, k2, k3]) assert_equal(body.filter(), [k1, f1, i1, i2, k2, k3])