diff options
Diffstat (limited to 'Lib/compiler')
| -rw-r--r-- | Lib/compiler/ast.py | 76 | ||||
| -rw-r--r-- | Lib/compiler/pycodegen.py | 89 | ||||
| -rw-r--r-- | Lib/compiler/symbols.py | 41 | ||||
| -rw-r--r-- | Lib/compiler/transformer.py | 53 |
4 files changed, 259 insertions, 0 deletions
diff --git a/Lib/compiler/ast.py b/Lib/compiler/ast.py index a10225b..064df50 100644 --- a/Lib/compiler/ast.py +++ b/Lib/compiler/ast.py @@ -1236,6 +1236,82 @@ class ListCompFor(Node): def __repr__(self): return "ListCompFor(%s, %s, %s)" % (repr(self.assign), repr(self.list), repr(self.ifs)) +class GenExpr(Node): + nodes["genexpr"] = "GenExpr" + def __init__(self, code): + self.code = code + self.argnames = ['[outmost-iterable]'] + self.varargs = self.kwargs = None + + def getChildren(self): + return self.code, + + def getChildNodes(self): + return self.code, + + def __repr__(self): + return "GenExpr(%s)" % (repr(self.code),) + +class GenExprInner(Node): + nodes["genexprinner"] = "GenExprInner" + def __init__(self, expr, quals): + self.expr = expr + self.quals = quals + + def getChildren(self): + children = [] + children.append(self.expr) + children.extend(flatten(self.quals)) + return tuple(children) + + def getChildNodes(self): + nodelist = [] + nodelist.append(self.expr) + nodelist.extend(flatten_nodes(self.quals)) + return tuple(nodelist) + + def __repr__(self): + return "GenExprInner(%s, %s)" % (repr(self.expr), repr(self.quals)) + +class GenExprFor(Node): + nodes["genexprfor"] = "GenExprFor" + def __init__(self, assign, iter, ifs): + self.assign = assign + self.iter = iter + self.ifs = ifs + self.is_outmost = False + + def getChildren(self): + children = [] + children.append(self.assign) + children.append(self.iter) + children.extend(flatten(self.ifs)) + return tuple(children) + + def getChildNodes(self): + nodelist = [] + nodelist.append(self.assign) + nodelist.append(self.iter) + nodelist.extend(flatten_nodes(self.ifs)) + return tuple(nodelist) + + def __repr__(self): + return "GenExprFor(%s, %s, %s)" % (repr(self.assign), repr(self.iter), repr(self.ifs)) + +class GenExprIf(Node): + nodes["genexprif"] = "GenExprIf" + def __init__(self, test): + self.test = test + + def getChildren(self): + return self.test, + + def getChildNodes(self): + return self.test, + + def __repr__(self): + return "GenExprIf(%s)" % (repr(self.test),) + klasses = globals() for k in nodes.keys(): nodes[k] = klasses[nodes[k]] diff --git a/Lib/compiler/pycodegen.py b/Lib/compiler/pycodegen.py index 292e859..ef5a0d3 100644 --- a/Lib/compiler/pycodegen.py +++ b/Lib/compiler/pycodegen.py @@ -619,6 +619,79 @@ class CodeGenerator: self.newBlock() self.emit('POP_TOP') + def visitGenExpr(self, node): + gen = GenExprCodeGenerator(node, self.scopes, self.class_name, + self.get_module()) + walk(node.code, gen) + gen.finish() + self.set_lineno(node) + frees = gen.scope.get_free_vars() + if frees: + for name in frees: + self.emit('LOAD_CLOSURE', name) + self.emit('LOAD_CONST', gen) + self.emit('MAKE_CLOSURE', 0) + else: + self.emit('LOAD_CONST', gen) + self.emit('MAKE_FUNCTION', 0) + + # precomputation of outmost iterable + self.visit(node.code.quals[0].iter) + self.emit('GET_ITER') + self.emit('CALL_FUNCTION', 1) + + def visitGenExprInner(self, node): + self.set_lineno(node) + # setup list + + stack = [] + for i, for_ in zip(range(len(node.quals)), node.quals): + start, anchor = self.visit(for_) + cont = None + for if_ in for_.ifs: + if cont is None: + cont = self.newBlock() + self.visit(if_, cont) + stack.insert(0, (start, cont, anchor)) + + self.visit(node.expr) + self.emit('YIELD_VALUE') + + for start, cont, anchor in stack: + if cont: + skip_one = self.newBlock() + self.emit('JUMP_FORWARD', skip_one) + self.startBlock(cont) + self.emit('POP_TOP') + self.nextBlock(skip_one) + self.emit('JUMP_ABSOLUTE', start) + self.startBlock(anchor) + self.emit('LOAD_CONST', None) + + def visitGenExprFor(self, node): + start = self.newBlock() + anchor = self.newBlock() + + if node.is_outmost: + self.loadName('[outmost-iterable]') + else: + self.visit(node.iter) + self.emit('GET_ITER') + + self.nextBlock(start) + self.set_lineno(node, force=True) + self.emit('FOR_ITER', anchor) + self.nextBlock() + self.visit(node.assign) + return start, anchor + + def visitGenExprIf(self, node, branch): + self.set_lineno(node, force=True) + self.visit(node.test) + self.emit('JUMP_IF_FALSE', branch) + self.newBlock() + self.emit('POP_TOP') + # exception related def visitAssert(self, node): @@ -1199,6 +1272,7 @@ class AbstractFunctionCode: klass.lambdaCount = klass.lambdaCount + 1 else: name = func.name + args, hasTupleArg = generateArgList(func.argnames) self.graph = pyassem.PyFlowGraph(name, func.filename, args, optimized=1) @@ -1263,6 +1337,21 @@ class FunctionCodeGenerator(NestedScopeMixin, AbstractFunctionCode, if self.scope.generator is not None: self.graph.setFlag(CO_GENERATOR) +class GenExprCodeGenerator(NestedScopeMixin, AbstractFunctionCode, + CodeGenerator): + super_init = CodeGenerator.__init__ # call be other init + scopes = None + + __super_init = AbstractFunctionCode.__init__ + + def __init__(self, gexp, scopes, class_name, mod): + self.scopes = scopes + self.scope = scopes[gexp] + self.__super_init(gexp, scopes, 1, class_name, mod) + self.graph.setFreeVars(self.scope.get_free_vars()) + self.graph.setCellVars(self.scope.get_cell_vars()) + self.graph.setFlag(CO_GENERATOR) + class AbstractClassCode: def __init__(self, klass, scopes, module): diff --git a/Lib/compiler/symbols.py b/Lib/compiler/symbols.py index fa668f1..6843e19 100644 --- a/Lib/compiler/symbols.py +++ b/Lib/compiler/symbols.py @@ -179,6 +179,21 @@ class ModuleScope(Scope): class FunctionScope(Scope): pass +class GenExprScope(Scope): + __super_init = Scope.__init__ + + __counter = 1 + + def __init__(self, module, klass=None): + i = self.__counter + self.__counter += 1 + self.__super_init("generator expression<%d>"%i, module, klass) + self.add_param('[outmost-iterable]') + + def get_names(self): + keys = Scope.get_names() + return keys + class LambdaScope(FunctionScope): __super_init = Scope.__init__ @@ -220,6 +235,32 @@ class SymbolVisitor: self.visit(node.code, scope) self.handle_free_vars(scope, parent) + def visitGenExpr(self, node, parent): + scope = GenExprScope(self.module, self.klass); + if parent.nested or isinstance(parent, FunctionScope) \ + or isinstance(parent, GenExprScope): + scope.nested = 1 + + self.scopes[node] = scope + self.visit(node.code, scope) + + self.handle_free_vars(scope, parent) + + def visitGenExprInner(self, node, scope): + for genfor in node.quals: + self.visit(genfor, scope) + + self.visit(node.expr, scope) + + def visitGenExprFor(self, node, scope): + self.visit(node.assign, scope, 1) + self.visit(node.iter, scope) + for if_ in node.ifs: + self.visit(if_, scope) + + def visitGenExprIf(self, node, scope): + self.visit(node.test, scope) + def visitLambda(self, node, parent, assign=0): # Lambda is an expression, so it could appear in an expression # context where assign is passed. The transformer should catch diff --git a/Lib/compiler/transformer.py b/Lib/compiler/transformer.py index 9074fc7..6832cf1 100644 --- a/Lib/compiler/transformer.py +++ b/Lib/compiler/transformer.py @@ -534,6 +534,12 @@ class Transformer: testlist1 = testlist exprlist = testlist + def testlist_gexp(self, nodelist): + if len(nodelist) == 2 and nodelist[1][0] == symbol.gen_for: + test = self.com_node(nodelist[0]) + return self.com_generator_expression(test, nodelist[1]) + return self.testlist(nodelist) + def test(self, nodelist): # and_test ('or' and_test)* | lambdef if len(nodelist) == 1 and nodelist[0][0] == symbol.lambdef: @@ -1085,6 +1091,48 @@ class Transformer: values.append(self.com_node(nodelist[i])) return List(values) + if hasattr(symbol, 'gen_for'): + def com_generator_expression(self, expr, node): + # gen_iter: gen_for | gen_if + # gen_for: 'for' exprlist 'in' test [gen_iter] + # gen_if: 'if' test [gen_iter] + + lineno = node[1][2] + fors = [] + while node: + t = node[1][1] + if t == 'for': + assignNode = self.com_assign(node[2], OP_ASSIGN) + genNode = self.com_node(node[4]) + newfor = GenExprFor(assignNode, genNode, []) + newfor.lineno = node[1][2] + fors.append(newfor) + if (len(node)) == 5: + node = None + else: + node = self.com_gen_iter(node[5]) + elif t == 'if': + test = self.com_node(node[2]) + newif = GenExprIf(test) + newif.lineno = node[1][2] + newfor.ifs.append(newif) + if len(node) == 3: + node = None + else: + node = self.com_gen_iter(node[3]) + else: + raise SyntaxError, \ + ("unexpected generator expression element: %s %d" + % (node, lineno)) + fors[0].is_outmost = True + n = GenExpr(GenExprInner(expr, fors)) + n.lineno = lineno + return n + + def com_gen_iter(self, node): + assert node[0] == symbol.gen_iter + return node[1] + def com_dictmaker(self, nodelist): # dictmaker: test ':' test (',' test ':' value)* [','] items = [] @@ -1122,6 +1170,8 @@ class Transformer: if node[0] == token.STAR or node[0] == token.DOUBLESTAR: break kw, result = self.com_argument(node, kw) + if len_nodelist != 2 and isinstance(result, GenExpr): + raise SyntaxError, 'generator expression needs parenthesis' args.append(result) else: # No broken by star arg, so skip the last one we processed. @@ -1148,6 +1198,9 @@ class Transformer: return CallFunc(primaryNode, args, star_node, dstar_node) def com_argument(self, nodelist, kw): + if len(nodelist) == 3 and nodelist[2][0] == symbol.gen_for: + test = self.com_node(nodelist[1]) + return 0, self.com_generator_expression(test, nodelist[2]) if len(nodelist) == 2: if kw: raise SyntaxError, "non-keyword arg after keyword arg" |