def call_Where(self, node: ast.Call, args: List[ast.AST]) -> ast.AST:
r'''
Transformation #1:
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))
Transformation #2:
seq.Select(x: f(x)).Where(y: g(y))
=> Where(Select(seq, x: f(x)), y: g(y))
Is turned into:
seq.Where(x: g(f(x))).Select(x: f(x))
=> Select(Where(seq, x: g(f(x)), f(x))
Transformation #3:
seq.SelectMany(x: f(x)).Where(y: g(y))
=> Where(SelectMany(seq, x: f(x)), y: g(y))
Is turned into:
seq.SelectMany(x: f(x).Where(y: g(y)))
=> SelectMany(seq, x: Where(f(x), g(y)))
'''
source = args[0]
filter = args[1]
assert isinstance(filter, ast.Lambda)
parent_where = self.visit(source)
if is_call_of(parent_where, 'Where'):
return self.visit_Where_of_Where(parent_where, filter)
elif is_call_of(parent_where, 'Select'):
return self.visit_Where_of_Select(parent_where, filter)
elif is_call_of(parent_where, 'SelectMany'):
return self.visit_Where_of_SelectMany(parent_where, filter)
else:
f = self.visit(filter)
if lambda_is_true(f):
return parent_where
else:
return function_call('Where', [parent_where, f])
def visit_Where(self, node):
r'''
Transformation #1:
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))
Transformation #2:
seq.Select(x: f(x)).Where(y: g(y))
=> Where(Select(seq, x: f(x)), y: g(y))
Is turned into:
seq.Where(x: g(f(x))).Select(x: f(x))
=> Select(Where(seq, x: g(f(x)), f(x))
Transformation #3:
seq.SelectMany(x: f(x)).Where(y: g(y))
=> Where(SelectMany(seq, x: f(x)), y: g(y))
Is turned into:
seq.SelectMany(x: f(x).Where(y: g(y)))
=> SelectMany(seq, x: Where(f(x), g(y)))
'''
parent_where = self.visit(node.source)
if type(parent_where) is Where:
return self.visit_Where_of_Where(parent_where, node.filter)
elif type(parent_where) is Select:
return self.visit_Where_of_Select(parent_where, node.filter)
elif type(parent_where) is SelectMany:
return self.visit_Where_of_SelectMany(parent_where, node.filter)
else:
f = self.visit(node.filter)
if lambda_is_true(f):
return parent_where
else:
return Where(parent_where, f)
def test_lambda_is_true_non_lambda():
assert lambda_is_true(ast.parse("True")) is False
def test_lambda_is_true_expression():
assert lambda_is_true(ast.parse("lambda x: x")) is False
def test_lambda_is_true_no():
assert lambda_is_true(ast.parse("lambda x: False")) is False
def test_lambda_is_true_yes():
assert lambda_is_true(ast.parse("lambda x: True")) is True
