def visit_Where_of_Where(self, parent, filter):
'''
seq.Where(x: f(x)).Where(x: g(x))
=> Where(Where(seq, x: f(x)), y: g(y))
is turned into
seq.Where(x: f(x) and g(y))
=> Where(seq, x: f(x) and g(y))
'''
func_f = parent.filter
func_g = filter
arg = arg_name()
return self.visit(
Where(
parent.source,
lambda_build(
arg,
ast.BoolOp(
ast.And(),
[lambda_call(arg, func_f),
lambda_call(arg, func_g)]))))
def visit_Where_of_Where(self, parent: ast.Call, filter: ast.Lambda):
'''
seq.Where(x: f(x)).Where(x: g(x))
=> Where(Where(seq, x: f(x)), y: g(y))
is turned into
seq.Where(x: f(x) and g(y))
=> Where(seq, x: f(x) and g(y))
'''
# Unpack arguments and f and g functions
_, args = unpack_Call(parent)
source = args[0]
func_f = args[1]
assert isinstance(func_f, ast.Lambda)
func_g = filter
arg = arg_name()
convolution = lambda_build(
arg,
ast.BoolOp(ast.And(),
[lambda_call(arg, func_f),
lambda_call(arg, func_g)])) # type: ast.AST
return self.visit(function_call('Where', [source, convolution]))
def test_call_wrap_single_arg():
ln = ast.parse("lambda x: x+1")
c = lambda_call("x", ln)
assert isinstance(c, ast.Call)
def test_call_wrap_list_arg():
ln = ast.parse("lambda x: x+1")
c = lambda_call(["x"], ln)
assert isinstance(c, ast.Call)
def test_call_wrap_single_arg():
l = ast.parse('lambda x: x+1')
c = lambda_call('x', l)
assert isinstance(c, ast.Call)
def test_call_wrap_list_arg():
l = ast.parse('lambda x: x+1')
c = lambda_call(['x'], l)
assert isinstance(c, ast.Call)
