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