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)