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])
