def visit_SelectMany_of_SelectMany(self, parent: ast.Call,
selection: ast.Lambda):
'''
Transformation #1:
seq.SelectMany(x: f(x)).SelectMany(y: f(y))
=> SelectMany(SelectMany(seq, x: f(x)), y: f(y))
is turned into:
seq.SelectMany(x: f(x).SelectMany(y: f(y)))
=> SelectMany(seq, x: SelectMany(f(x), y: f(y)))
'''
_, args = unpack_Call(parent)
assert (args is not None) and len(args) == 2
seq = args[0]
func_f = args[1]
assert isinstance(func_f, ast.Lambda)
func_g = selection
captured_arg = func_f.args.args[0].arg
captured_body = func_f.body
new_select = function_call(
'SelectMany',
[cast(ast.AST, captured_body),
cast(ast.AST, func_g)])
new_select_lambda = lambda_build(captured_arg, new_select)
new_selectmany = function_call(
'SelectMany', [seq, cast(ast.AST, new_select_lambda)])
return new_selectmany
def select_method_call_on_first(self, node: ast.Call):
"""Turn
First(seq).method(args)
into
First(Select(seq, s: s.method(args)))
"""
# Extract the call info
assert isinstance(node.func, ast.Attribute)
method_name = node.func.attr
method_args = node.args
method_keywords = (
node.keywords
if hasattr(node, "keywords") else node.kwargs # type: ignore
)
assert isinstance(node.func.value, ast.Call)
seq = node.func.value.args[0]
# Now rebuild the call
a = arg_name()
call_args = {
"func": ast.Attribute(value=ast.Name(a, ast.Load()),
attr=method_name),
"args": method_args,
}
if hasattr(node, "keywords"):
call_args["keywords"] = method_keywords
else:
call_args["kwargs"] = method_keywords
seq_a_call = ast.Call(**call_args)
select = make_Select(seq, lambda_build(a, seq_a_call))
return self.visit(function_call("First", [cast(ast.AST, select)]))
def call_SelectMany(self, node: ast.Call, args: List[ast.AST]):
r'''
Transformation #1:
seq.SelectMany(x: f(x)).SelectMany(y: f(y))
=> SelectMany(SelectMany(seq, x: f(x)), y: f(y))
is turned into:
seq.SelectMany(x: f(x).SelectMany(y: f(y)))
=> SelectMany(seq, x: SelectMany(f(x), y: f(y)))
Transformation #2:
seq.Select(x: f(x)).SelectMany(y: g(y))
=> SelectMany(Select(seq, x: f(x)), y:g(y))
is turned into
seq.SelectMany(x: g(f(x)))
=> SelectMany(seq, x: g(f(x)))
Transformation #3:
seq.Where(x: f(x)).SelectMany(y: g(y))
'''
selection = args[1]
assert isinstance(selection, ast.Lambda)
parent_select = self.visit(args[0])
if is_call_of(parent_select, 'SelectMany'):
return self.visit_SelectMany_of_SelectMany(parent_select,
selection)
elif is_call_of(parent_select, 'Select'):
return self.visit_SelectMany_of_Select(parent_select, selection)
else:
return function_call(
'SelectMany',
[parent_select, self.visit(selection)])
def visit_Where_of_SelectMany(self, parent, filter):
'''
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)))
'''
_, args = unpack_Call(parent)
seq = args[0]
func_f = args[1]
assert isinstance(func_f, ast.Lambda)
func_g = filter
lambda_where = lambda_body_replace(
func_f, function_call("Where", [lambda_body(func_f), func_g]))
return self.visit(function_call('SelectMany', [seq, lambda_where]))
def visit_Call(self, call_node: ast.Call) -> Optional[ast.AST]:
node = self.generic_visit(call_node)
if node is None or not isinstance(node, ast.Call):
return node
if not isinstance(node.func, ast.Attribute):
return node
if node.func.attr not in function_names:
return node
return function_call(node.func.attr, cast(List[ast.AST], [node.func.value] + node.args))
def test_function_call_simple():
a = function_call('dude', [as_ast(1)])
print(ast.dump(ast.parse('dude(1)')))
if sys.version_info < (3, 8):
expected = "Call(func=Name(id='dude', ctx=Load()), args=[Num(n=1)], keywords=[])"
elif sys.version_info < (3, 9):
expected = "Call(func=Name(id='dude', ctx=Load()), args=[Constant(value=1, kind=None)], keywords=[])"
else:
expected = "Call(func=Name(id='dude', ctx=Load()), args=[Constant(value=1)], keywords=[])"
assert expected == ast.dump(a)
def _generate_count_call(seq: ast.AST, lambda_string: str = "lambda acc,v: acc+1") -> ast.Call:
r"""
Given a sequence, generate an Aggregate call that will count the number
of items in the sequence.
seq: The sequence to be counted
returns:
agg_ast - An ast call to the Aggregate call.
"""
agg_lambda = cast(ast.Expr, ast.parse(lambda_string).body[0]).value
agg_start = ast.Num(0) if sys.version_info < (3, 8, 0) else ast.Constant(0, kind=None)
return function_call("Aggregate", [seq, cast(ast.AST, agg_start), cast(ast.AST, agg_lambda)])
def visit_Attribute_Of_First(self, first: ast.AST, attr: str):
"""
Convert a seq.First().attr
==>
seq.Select(l: l.attr).First()
Other work will do the conversion as needed.
"""
# Build the select that starts from the source and does the slice.
a = arg_name()
select = make_Select(
first,
lambda_build(
a, ast.Attribute(value=ast.Name(a, ast.Load()), attr=attr)))
return self.visit(function_call("First", [cast(ast.AST, select)]))
def visit_Select_of_SelectMany(self, parent, selection):
r'''
seq.SelectMany(x: f(x)).Select(y: g(y))
=> Select(SelectMany(seq, x: f(x)), y: g(y))
is turned into
seq.SelectMany(x: f(x).Select(y: g(y)))
=> SelectMany(seq, x: Select(f(x), y: g(y)))
'''
(_, args) = unpack_Call(parent)
source = args[0]
func_f = args[1]
assert isinstance(func_f, ast.Lambda)
func_g = selection
lambda_select = \
lambda_body_replace(func_f, make_Select(lambda_body(func_f), func_g)) # type: ast.AST
return self.visit(function_call('SelectMany', [source, lambda_select]))
def visit_Subscript_Of_First(self, first: ast.AST, s):
'''
Convert a seq.First()[0]
==>
seq.Select(l: l[0]).First()
Other work will do the conversion as needed.
'''
# Build the select that starts from the source and does the slice.
a = arg_name()
select = make_Select(
first,
lambda_build(a,
ast.Subscript(ast.Name(a, ast.Load()), s,
ast.Load())))
return self.visit(function_call('First', [cast(ast.AST, select)]))
def _generate_count_call(seq: ast.AST,
lambda_string: str = "lambda acc,v: acc+1"
) -> ast.Call:
r'''
Given a sequence, generate an Aggregate call that will count the number
of items in the sequence.
seq: The sequence to be counted
returns:
agg_ast - An ast call to the Aggregate call.
'''
agg_lambda = cast(ast.Expr, ast.parse(lambda_string).body[0]).value
agg_start = ast.Num(0)
return function_call(
'Aggregate',
[seq, cast(ast.AST, agg_start),
cast(ast.AST, agg_lambda)])
def visit_SelectMany_of_Select(self, parent_select: ast.Call,
selection: ast.Lambda):
'''
seq.Select(x: f(x)).SelectMany(y: g(y))
=> SelectMany(Select(seq, x: f(x)), y:g(y))
is turned into
seq.SelectMany(x: f(x).Select(y: g(y)))
=> SelectMany(seq, x: Select(f(x), y: g(y)))
'''
_, select_args = unpack_Call(parent_select)
assert (select_args is not None) and len(select_args) == 2
seq = select_args[0]
func_f = select_args[1]
assert isinstance(func_f, ast.Lambda)
func_g = selection
w = function_call('SelectMany',
[seq, self.visit(convolute(func_g, func_f))])
return w
def visit_Where_of_Select(self, parent, filter):
'''
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))
'''
_, args = unpack_Call(parent)
source = args[0]
func_f = args[1]
assert isinstance(func_f, ast.Lambda)
func_g = filter
w = function_call(
'Where', [source, self.visit(convolute(func_g, func_f))])
s = make_Select(w, func_f)
# Recursively visit this mess to see if the Where needs to move further up.
return self.visit(s)
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_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_good_ev_call():
'Test a good call for an event collection'
doit = function_call('Jets', [ast.parse('"antikt"').body[0].value]) # type: ignore
getCollection(collections[0], doit)
def test_good_ev_call_arg_number_bad():
'Wrong number of arguments to a collection call'
doit = function_call('Jets', [ast.parse('"antikt"').body[0].value, ast.parse('55').body[0].value]) # type: ignore
with pytest.raises(ValueError):
getCollection(collections[0], doit)
def test_good_ev_call_bad_arg_type():
'Bad collection type'
doit = function_call('Jets', [ast.parse('55').body[0].value]) # type: ignore
with pytest.raises(ValueError):
getCollection(collections[0], doit)
def make_Select(source: ast.AST, selection: ast.AST):
'Make a select, and return source is selection is an identity'
return source if lambda_is_identity(selection) else function_call(
'Select', [source, selection])
def test_function_call_simple():
a = function_call('dude', [as_ast(1)])
print(ast.dump(ast.parse('dude(1)')))
expected = "Call(func=Name(id='dude', ctx=Load()), args=[Num(n=1)], keywords=[])"
assert expected == ast.dump(a)
