def make_Iterator(self, gen):
if gen.ifs:
ldFilter = ast.Lambda(
ast.arguments([ast.Name(gen.target.id, ast.Param(), None)],
None, [], [], None, []),
ast.BoolOp(ast.And(), gen.ifs)
if len(gen.ifs) > 1 else gen.ifs[0])
ifilterName = ast.Attribute(value=ast.Name(id=MODULE,
ctx=ast.Load(),
annotation=None),
attr=IFILTER,
ctx=ast.Load())
return ast.Call(ifilterName, [ldFilter, gen.iter], [])
else:
return gen.iter
def visit_ListComp(self, node):
if node in self.optimizable_comprehension:
self.update = True
self.generic_visit(node)
iterList = []
varList = []
for gen in node.generators:
iterList.append(self.make_Iterator(gen))
varList.append(ast.Name(gen.target.id, ast.Param(), None))
# If dim = 1, product is useless
if len(iterList) == 1:
iterAST = iterList[0]
varAST = ast.arguments([varList[0]], None, [], [], None, [])
else:
self.use_itertools = True
prodName = ast.Attribute(value=ast.Name(id=mangle('itertools'),
ctx=ast.Load(),
annotation=None),
attr='product',
ctx=ast.Load())
varid = varList[0].id # retarget this id, it's free
renamings = {v.id: (i, ) for i, v in enumerate(varList)}
node.elt = ConvertToTuple(varid, renamings).visit(node.elt)
iterAST = ast.Call(prodName, iterList, [])
varAST = ast.arguments([ast.Name(varid, ast.Param(), None)],
None, [], [], None, [])
mapName = ast.Attribute(value=ast.Name(id='__builtin__',
ctx=ast.Load(),
annotation=None),
attr='map',
ctx=ast.Load())
ldBodymap = node.elt
ldmap = ast.Lambda(varAST, ldBodymap)
return ast.Call(mapName, [ldmap, iterAST], [])
else:
return self.generic_visit(node)
def create_lambda_node(func_or_expr_node, is_if_expr=False):
body = func_or_expr_node
if not is_if_expr:
body = gast.Call(func=gast.Name(id=func_or_expr_node.name,
ctx=gast.Load(),
annotation=None,
type_comment=None),
args=[func_or_expr_node.args],
keywords=[])
lambda_node = gast.Lambda(args=gast.arguments(args=[],
posonlyargs=[],
vararg=None,
kwonlyargs=[],
kw_defaults=None,
kwarg=None,
defaults=[]),
body=body)
return lambda_node
def visitComp(self, node, make_attr):
if node in self.optimizable_comprehension:
self.update = True
self.generic_visit(node)
iters = [self.make_Iterator(gen) for gen in node.generators]
variables = [
ast.Name(gen.target.id, ast.Param(), None, None)
for gen in node.generators
]
# If dim = 1, product is useless
if len(iters) == 1:
iterAST = iters[0]
varAST = ast.arguments([variables[0]], [], None, [], [], None,
[])
else:
self.use_itertools = True
prodName = ast.Attribute(value=ast.Name(id=mangle('itertools'),
ctx=ast.Load(),
annotation=None,
type_comment=None),
attr='product',
ctx=ast.Load())
varid = variables[0].id # retarget this id, it's free
renamings = {v.id: (i, ) for i, v in enumerate(variables)}
node.elt = ConvertToTuple(varid, renamings).visit(node.elt)
iterAST = ast.Call(prodName, iters, [])
varAST = ast.arguments(
[ast.Name(varid, ast.Param(), None, None)], [], None, [],
[], None, [])
ldBodymap = node.elt
ldmap = ast.Lambda(varAST, ldBodymap)
return make_attr(ldmap, iterAST)
else:
return self.generic_visit(node)
def visit_GeneratorExp(self, node):
if node in self.optimizable_comprehension:
self.update = True
self.generic_visit(node)
iters = [self.make_Iterator(gen) for gen in node.generators]
variables = [
ast.Name(gen.target.id, ast.Param(), None)
for gen in node.generators
]
# If dim = 1, product is useless
if len(iters) == 1:
iterAST = iters[0]
varAST = ast.arguments([variables[0]], None, [], [], None, [])
else:
prodName = ast.Attribute(value=ast.Name(id='itertools',
ctx=ast.Load(),
annotation=None),
attr='product',
ctx=ast.Load())
iterAST = ast.Call(prodName, iters, [])
varAST = ast.arguments([ast.Tuple(variables, ast.Store())],
None, [], [], None, [])
imapName = ast.Attribute(value=ast.Name(id=MODULE,
ctx=ast.Load(),
annotation=None),
attr=IMAP,
ctx=ast.Load())
ldBodyimap = node.elt
ldimap = ast.Lambda(varAST, ldBodyimap)
return ast.Call(imapName, [ldimap, iterAST], [])
else:
return self.generic_visit(node)
