pycodegen.py revision 2e4cc7e0d84747826d3d2546d1bccd9c40a455c2 
1import imp
2import os
3import marshal
4import stat
5import string
6import struct
7import sys
8import types
9from cStringIO import StringIO
10
11from compiler import ast, parse, walk, syntax
12from compiler import pyassem, misc, future, symbols
13from compiler.consts import SC_LOCAL, SC_GLOBAL, SC_FREE, SC_CELL
14from compiler.consts import CO_VARARGS, CO_VARKEYWORDS, CO_NEWLOCALS,\
15     CO_NESTED, CO_GENERATOR, CO_GENERATOR_ALLOWED, CO_FUTURE_DIVISION
16from compiler.pyassem import TupleArg
17
18# Do we have Python 1.x or Python 2.x?
19try:
20    VERSION = sys.version_info[0]
21except AttributeError:
22    VERSION = 1
23
24callfunc_opcode_info = {
25    # (Have *args, Have **args) : opcode
26    (0,0) : "CALL_FUNCTION",
27    (1,0) : "CALL_FUNCTION_VAR",
28    (0,1) : "CALL_FUNCTION_KW",
29    (1,1) : "CALL_FUNCTION_VAR_KW",
30}
31
32LOOP = 1
33EXCEPT = 2
34TRY_FINALLY = 3
35END_FINALLY = 4
36
37class BlockStack(misc.Stack):
38    __super_init = misc.Stack.__init__
39
40    def __init__(self):
41        self.__super_init(self)
42        self.loop = None
43
44def compile(filename, display=0):
45    f = open(filename)
46    buf = f.read()
47    f.close()
48    mod = Module(buf, filename)
49    try:
50        mod.compile(display)
51    except SyntaxError, err:
52        raise
53    else:
54        f = open(filename + "c", "wb")
55        mod.dump(f)
56        f.close()
57
58class Module:
59    def __init__(self, source, filename):
60        self.filename = os.path.abspath(filename)
61        self.source = source
62        self.code = None
63
64    def compile(self, display=0):
65        tree = parse(self.source)
66        misc.set_filename(self.filename, tree)
67        syntax.check(tree)
68        gen = ModuleCodeGenerator(tree)
69        if display:
70            import pprint
71            print pprint.pprint(tree)
72        self.code = gen.getCode()
73
74    def dump(self, f):
75        f.write(self.getPycHeader())
76        marshal.dump(self.code, f)
77
78    MAGIC = imp.get_magic()
79
80    def getPycHeader(self):
81        # compile.c uses marshal to write a long directly, with
82        # calling the interface that would also generate a 1-byte code
83        # to indicate the type of the value.  simplest way to get the
84        # same effect is to call marshal and then skip the code.
85        mtime = os.stat(self.filename)[stat.ST_MTIME]
86        mtime = struct.pack('i', mtime)
87        return self.MAGIC + mtime
88
89class LocalNameFinder:
90    """Find local names in scope"""
91    def __init__(self, names=()):
92        self.names = misc.Set()
93        self.globals = misc.Set()
94        for name in names:
95            self.names.add(name)
96
97    # XXX list comprehensions and for loops
98
99    def getLocals(self):
100        for elt in self.globals.elements():
101            if self.names.has_elt(elt):
102                self.names.remove(elt)
103        return self.names
104
105    def visitDict(self, node):
106        pass
107
108    def visitGlobal(self, node):
109        for name in node.names:
110            self.globals.add(name)
111
112    def visitFunction(self, node):
113        self.names.add(node.name)
114
115    def visitLambda(self, node):
116        pass
117
118    def visitImport(self, node):
119        for name, alias in node.names:
120            self.names.add(alias or name)
121
122    def visitFrom(self, node):
123        for name, alias in node.names:
124            self.names.add(alias or name)
125
126    def visitClass(self, node):
127        self.names.add(node.name)
128
129    def visitAssName(self, node):
130        self.names.add(node.name)
131
132def is_constant_false(node):
133    if isinstance(node, ast.Const):
134        if not node.value:
135            return 1
136    return 0
137
138class CodeGenerator:
139    """Defines basic code generator for Python bytecode
140
141    This class is an abstract base class.  Concrete subclasses must
142    define an __init__() that defines self.graph and then calls the
143    __init__() defined in this class.
144
145    The concrete class must also define the class attributes
146    NameFinder, FunctionGen, and ClassGen.  These attributes can be
147    defined in the initClass() method, which is a hook for
148    initializing these methods after all the classes have been
149    defined.
150    """
151
152    optimized = 0 # is namespace access optimized?
153    __initialized = None
154    class_name = None # provide default for instance variable
155
156    def __init__(self):
157        if self.__initialized is None:
158            self.initClass()
159            self.__class__.__initialized = 1
160        self.checkClass()
161        self.locals = misc.Stack()
162        self.setups = misc.Stack()
163        self.curStack = 0
164        self.maxStack = 0
165        self.last_lineno = None
166        self._setupGraphDelegation()
167        self._div_op = "BINARY_DIVIDE"
168
169        # XXX set flags based on future features
170        futures = self.get_module().futures
171        for feature in futures:
172            if feature == "division":
173                self.graph.setFlag(CO_FUTURE_DIVISION)
174                self._div_op = "BINARY_TRUE_DIVIDE"
175            elif feature == "generators":
176                self.graph.setFlag(CO_GENERATOR_ALLOWED)
177
178    def initClass(self):
179        """This method is called once for each class"""
180
181    def checkClass(self):
182        """Verify that class is constructed correctly"""
183        try:
184            assert hasattr(self, 'graph')
185            assert getattr(self, 'NameFinder')
186            assert getattr(self, 'FunctionGen')
187            assert getattr(self, 'ClassGen')
188        except AssertionError, msg:
189            intro = "Bad class construction for %s" % self.__class__.__name__
190            raise AssertionError, intro
191
192    def _setupGraphDelegation(self):
193        self.emit = self.graph.emit
194        self.newBlock = self.graph.newBlock
195        self.startBlock = self.graph.startBlock
196        self.nextBlock = self.graph.nextBlock
197        self.setDocstring = self.graph.setDocstring
198
199    def getCode(self):
200        """Return a code object"""
201        return self.graph.getCode()
202
203    def mangle(self, name):
204        if self.class_name is not None:
205            return misc.mangle(name, self.class_name)
206        else:
207            return name
208
209    def parseSymbols(self, tree):
210        s = symbols.SymbolVisitor()
211        walk(tree, s)
212        return s.scopes
213
214    def get_module(self):
215        raise RuntimeError, "should be implemented by subclasses"
216
217    # Next five methods handle name access
218
219    def isLocalName(self, name):
220        return self.locals.top().has_elt(name)
221
222    def storeName(self, name):
223        self._nameOp('STORE', name)
224
225    def loadName(self, name):
226        self._nameOp('LOAD', name)
227
228    def delName(self, name):
229        self._nameOp('DELETE', name)
230
231    def _nameOp(self, prefix, name):
232        name = self.mangle(name)
233        scope = self.scope.check_name(name)
234        if scope == SC_LOCAL:
235            if not self.optimized:
236                self.emit(prefix + '_NAME', name)
237            else:
238                self.emit(prefix + '_FAST', name)
239        elif scope == SC_GLOBAL:
240            if not self.optimized:
241                self.emit(prefix + '_NAME', name)
242            else:
243                self.emit(prefix + '_GLOBAL', name)
244        elif scope == SC_FREE or scope == SC_CELL:
245            self.emit(prefix + '_DEREF', name)
246        else:
247            raise RuntimeError, "unsupported scope for var %s: %d" % \
248                  (name, scope)
249
250    def _implicitNameOp(self, prefix, name):
251        """Emit name ops for names generated implicitly by for loops
252
253        The interpreter generates names that start with a period or
254        dollar sign.  The symbol table ignores these names because
255        they aren't present in the program text.
256        """
257        if self.optimized:
258            self.emit(prefix + '_FAST', name)
259        else:
260            self.emit(prefix + '_NAME', name)
261
262    def set_lineno(self, node, force=0):
263        """Emit SET_LINENO if node has lineno attribute and it is
264        different than the last lineno emitted.
265
266        Returns true if SET_LINENO was emitted.
267
268        There are no rules for when an AST node should have a lineno
269        attribute.  The transformer and AST code need to be reviewed
270        and a consistent policy implemented and documented.  Until
271        then, this method works around missing line numbers.
272        """
273        lineno = getattr(node, 'lineno', None)
274        if lineno is not None and (lineno != self.last_lineno
275                                   or force):
276            self.emit('SET_LINENO', lineno)
277            self.last_lineno = lineno
278            return 1
279        return 0
280
281    # The first few visitor methods handle nodes that generator new
282    # code objects.  They use class attributes to determine what
283    # specialized code generators to use.
284
285    NameFinder = LocalNameFinder
286    FunctionGen = None
287    ClassGen = None
288
289    def visitModule(self, node):
290        self.scopes = self.parseSymbols(node)
291        self.scope = self.scopes[node]
292        self.emit('SET_LINENO', 0)
293        if node.doc:
294            self.emit('LOAD_CONST', node.doc)
295            self.storeName('__doc__')
296        lnf = walk(node.node, self.NameFinder(), verbose=0)
297        self.locals.push(lnf.getLocals())
298        self.visit(node.node)
299        self.emit('LOAD_CONST', None)
300        self.emit('RETURN_VALUE')
301
302    def visitFunction(self, node):
303        self._visitFuncOrLambda(node, isLambda=0)
304        if node.doc:
305            self.setDocstring(node.doc)
306        self.storeName(node.name)
307
308    def visitLambda(self, node):
309        self._visitFuncOrLambda(node, isLambda=1)
310
311    def _visitFuncOrLambda(self, node, isLambda=0):
312        gen = self.FunctionGen(node, self.scopes, isLambda,
313                               self.class_name, self.get_module())
314        walk(node.code, gen)
315        gen.finish()
316        self.set_lineno(node)
317        for default in node.defaults:
318            self.visit(default)
319        frees = gen.scope.get_free_vars()
320        if frees:
321            for name in frees:
322                self.emit('LOAD_CLOSURE', name)
323            self.emit('LOAD_CONST', gen)
324            self.emit('MAKE_CLOSURE', len(node.defaults))
325        else:
326            self.emit('LOAD_CONST', gen)
327            self.emit('MAKE_FUNCTION', len(node.defaults))
328
329    def visitClass(self, node):
330        gen = self.ClassGen(node, self.scopes,
331                            self.get_module())
332        if node.doc:
333            self.emit('LOAD_CONST', node.doc)
334            self.storeName('__doc__')
335        walk(node.code, gen)
336        gen.finish()
337        self.set_lineno(node)
338        self.emit('LOAD_CONST', node.name)
339        for base in node.bases:
340            self.visit(base)
341        self.emit('BUILD_TUPLE', len(node.bases))
342        frees = gen.scope.get_free_vars()
343        for name in frees:
344            self.emit('LOAD_CLOSURE', name)
345        self.emit('LOAD_CONST', gen)
346        if frees:
347            self.emit('MAKE_CLOSURE', 0)
348        else:
349            self.emit('MAKE_FUNCTION', 0)
350        self.emit('CALL_FUNCTION', 0)
351        self.emit('BUILD_CLASS')
352        self.storeName(node.name)
353
354    # The rest are standard visitor methods
355
356    # The next few implement control-flow statements
357
358    def visitIf(self, node):
359        end = self.newBlock()
360        numtests = len(node.tests)
361        for i in range(numtests):
362            test, suite = node.tests[i]
363            if is_constant_false(test):
364                # XXX will need to check generator stuff here
365                continue
366            self.set_lineno(test)
367            self.visit(test)
368            nextTest = self.newBlock()
369            self.emit('JUMP_IF_FALSE', nextTest)
370            self.nextBlock()
371            self.emit('POP_TOP')
372            self.visit(suite)
373            self.emit('JUMP_FORWARD', end)
374            self.startBlock(nextTest)
375            self.emit('POP_TOP')
376        if node.else_:
377            self.visit(node.else_)
378        self.nextBlock(end)
379
380    def visitWhile(self, node):
381        self.set_lineno(node)
382
383        loop = self.newBlock()
384        else_ = self.newBlock()
385
386        after = self.newBlock()
387        self.emit('SETUP_LOOP', after)
388
389        self.nextBlock(loop)
390        self.setups.push((LOOP, loop))
391
392        self.set_lineno(node, force=1)
393        self.visit(node.test)
394        self.emit('JUMP_IF_FALSE', else_ or after)
395
396        self.nextBlock()
397        self.emit('POP_TOP')
398        self.visit(node.body)
399        self.emit('JUMP_ABSOLUTE', loop)
400
401        self.startBlock(else_) # or just the POPs if not else clause
402        self.emit('POP_TOP')
403        self.emit('POP_BLOCK')
404        self.setups.pop()
405        if node.else_:
406            self.visit(node.else_)
407        self.nextBlock(after)
408
409    def visitFor(self, node):
410        start = self.newBlock()
411        anchor = self.newBlock()
412        after = self.newBlock()
413        self.setups.push((LOOP, start))
414
415        self.set_lineno(node)
416        self.emit('SETUP_LOOP', after)
417        self.visit(node.list)
418        self.emit('GET_ITER')
419
420        self.nextBlock(start)
421        self.set_lineno(node, force=1)
422        self.emit('FOR_ITER', anchor)
423        self.visit(node.assign)
424        self.visit(node.body)
425        self.emit('JUMP_ABSOLUTE', start)
426        self.nextBlock(anchor)
427        self.emit('POP_BLOCK')
428        self.setups.pop()
429        if node.else_:
430            self.visit(node.else_)
431        self.nextBlock(after)
432
433    def visitBreak(self, node):
434        if not self.setups:
435            raise SyntaxError, "'break' outside loop (%s, %d)" % \
436                  (node.filename, node.lineno)
437        self.set_lineno(node)
438        self.emit('BREAK_LOOP')
439
440    def visitContinue(self, node):
441        if not self.setups:
442            raise SyntaxError, "'continue' outside loop (%s, %d)" % \
443                  (node.filename, node.lineno)
444        kind, block = self.setups.top()
445        if kind == LOOP:
446            self.set_lineno(node)
447            self.emit('JUMP_ABSOLUTE', block)
448            self.nextBlock()
449        elif kind == EXCEPT or kind == TRY_FINALLY:
450            self.set_lineno(node)
451            # find the block that starts the loop
452            top = len(self.setups)
453            while top > 0:
454                top = top - 1
455                kind, loop_block = self.setups[top]
456                if kind == LOOP:
457                    break
458            if kind != LOOP:
459                raise SyntaxError, "'continue' outside loop (%s, %d)" % \
460                      (node.filename, node.lineno)
461            self.emit('CONTINUE_LOOP', loop_block)
462            self.nextBlock()
463        elif kind == END_FINALLY:
464            msg = "'continue' not allowed inside 'finally' clause (%s, %d)"
465            raise SyntaxError, msg % (node.filename, node.lineno)
466
467    def visitTest(self, node, jump):
468        end = self.newBlock()
469        for child in node.nodes[:-1]:
470            self.visit(child)
471            self.emit(jump, end)
472            self.nextBlock()
473            self.emit('POP_TOP')
474        self.visit(node.nodes[-1])
475        self.nextBlock(end)
476
477    def visitAnd(self, node):
478        self.visitTest(node, 'JUMP_IF_FALSE')
479
480    def visitOr(self, node):
481        self.visitTest(node, 'JUMP_IF_TRUE')
482
483    def visitCompare(self, node):
484        self.visit(node.expr)
485        cleanup = self.newBlock()
486        for op, code in node.ops[:-1]:
487            self.visit(code)
488            self.emit('DUP_TOP')
489            self.emit('ROT_THREE')
490            self.emit('COMPARE_OP', op)
491            self.emit('JUMP_IF_FALSE', cleanup)
492            self.nextBlock()
493            self.emit('POP_TOP')
494        # now do the last comparison
495        if node.ops:
496            op, code = node.ops[-1]
497            self.visit(code)
498            self.emit('COMPARE_OP', op)
499        if len(node.ops) > 1:
500            end = self.newBlock()
501            self.emit('JUMP_FORWARD', end)
502            self.startBlock(cleanup)
503            self.emit('ROT_TWO')
504            self.emit('POP_TOP')
505            self.nextBlock(end)
506
507    # list comprehensions
508    __list_count = 0
509
510    def visitListComp(self, node):
511        self.set_lineno(node)
512        # setup list
513        append = "$append%d" % self.__list_count
514        self.__list_count = self.__list_count + 1
515        self.emit('BUILD_LIST', 0)
516        self.emit('DUP_TOP')
517        self.emit('LOAD_ATTR', 'append')
518        self._implicitNameOp('STORE', append)
519
520        stack = []
521        for i, for_ in zip(range(len(node.quals)), node.quals):
522            start, anchor = self.visit(for_)
523            cont = None
524            for if_ in for_.ifs:
525                if cont is None:
526                    cont = self.newBlock()
527                self.visit(if_, cont)
528            stack.insert(0, (start, cont, anchor))
529
530        self._implicitNameOp('LOAD', append)
531        self.visit(node.expr)
532        self.emit('CALL_FUNCTION', 1)
533        self.emit('POP_TOP')
534
535        for start, cont, anchor in stack:
536            if cont:
537                skip_one = self.newBlock()
538                self.emit('JUMP_FORWARD', skip_one)
539                self.startBlock(cont)
540                self.emit('POP_TOP')
541                self.nextBlock(skip_one)
542            self.emit('JUMP_ABSOLUTE', start)
543            self.startBlock(anchor)
544        self._implicitNameOp('DELETE', append)
545
546        self.__list_count = self.__list_count - 1
547
548    def visitListCompFor(self, node):
549        start = self.newBlock()
550        anchor = self.newBlock()
551
552        self.visit(node.list)
553        self.emit('GET_ITER')
554        self.nextBlock(start)
555        self.emit('SET_LINENO', node.lineno)
556        self.emit('FOR_ITER', anchor)
557        self.nextBlock()
558        self.visit(node.assign)
559        return start, anchor
560
561    def visitListCompIf(self, node, branch):
562        self.set_lineno(node, force=1)
563        self.visit(node.test)
564        self.emit('JUMP_IF_FALSE', branch)
565        self.newBlock()
566        self.emit('POP_TOP')
567
568    # exception related
569
570    def visitAssert(self, node):
571        # XXX would be interesting to implement this via a
572        # transformation of the AST before this stage
573        end = self.newBlock()
574        self.set_lineno(node)
575        # XXX __debug__ and AssertionError appear to be special cases
576        # -- they are always loaded as globals even if there are local
577        # names.  I guess this is a sort of renaming op.
578        self.emit('LOAD_GLOBAL', '__debug__')
579        self.emit('JUMP_IF_FALSE', end)
580        self.nextBlock()
581        self.emit('POP_TOP')
582        self.visit(node.test)
583        self.emit('JUMP_IF_TRUE', end)
584        self.nextBlock()
585        self.emit('POP_TOP')
586        self.emit('LOAD_GLOBAL', 'AssertionError')
587        if node.fail:
588            self.visit(node.fail)
589            self.emit('RAISE_VARARGS', 2)
590        else:
591            self.emit('RAISE_VARARGS', 1)
592        self.nextBlock(end)
593        self.emit('POP_TOP')
594
595    def visitRaise(self, node):
596        self.set_lineno(node)
597        n = 0
598        if node.expr1:
599            self.visit(node.expr1)
600            n = n + 1
601        if node.expr2:
602            self.visit(node.expr2)
603            n = n + 1
604        if node.expr3:
605            self.visit(node.expr3)
606            n = n + 1
607        self.emit('RAISE_VARARGS', n)
608
609    def visitTryExcept(self, node):
610        body = self.newBlock()
611        handlers = self.newBlock()
612        end = self.newBlock()
613        if node.else_:
614            lElse = self.newBlock()
615        else:
616            lElse = end
617        self.set_lineno(node)
618        self.emit('SETUP_EXCEPT', handlers)
619        self.nextBlock(body)
620        self.setups.push((EXCEPT, body))
621        self.visit(node.body)
622        self.emit('POP_BLOCK')
623        self.setups.pop()
624        self.emit('JUMP_FORWARD', lElse)
625        self.startBlock(handlers)
626
627        last = len(node.handlers) - 1
628        for i in range(len(node.handlers)):
629            expr, target, body = node.handlers[i]
630            self.set_lineno(expr)
631            if expr:
632                self.emit('DUP_TOP')
633                self.visit(expr)
634                self.emit('COMPARE_OP', 'exception match')
635                next = self.newBlock()
636                self.emit('JUMP_IF_FALSE', next)
637                self.nextBlock()
638                self.emit('POP_TOP')
639            self.emit('POP_TOP')
640            if target:
641                self.visit(target)
642            else:
643                self.emit('POP_TOP')
644            self.emit('POP_TOP')
645            self.visit(body)
646            self.emit('JUMP_FORWARD', end)
647            if expr:
648                self.nextBlock(next)
649            else:
650                self.nextBlock()
651            if expr: # XXX
652                self.emit('POP_TOP')
653        self.emit('END_FINALLY')
654        if node.else_:
655            self.nextBlock(lElse)
656            self.visit(node.else_)
657        self.nextBlock(end)
658
659    def visitTryFinally(self, node):
660        body = self.newBlock()
661        final = self.newBlock()
662        self.set_lineno(node)
663        self.emit('SETUP_FINALLY', final)
664        self.nextBlock(body)
665        self.setups.push((TRY_FINALLY, body))
666        self.visit(node.body)
667        self.emit('POP_BLOCK')
668        self.setups.pop()
669        self.emit('LOAD_CONST', None)
670        self.nextBlock(final)
671        self.setups.push((END_FINALLY, final))
672        self.visit(node.final)
673        self.emit('END_FINALLY')
674        self.setups.pop()
675
676    # misc
677
678    def visitDiscard(self, node):
679        self.set_lineno(node)
680        self.visit(node.expr)
681        self.emit('POP_TOP')
682
683    def visitConst(self, node):
684        self.emit('LOAD_CONST', node.value)
685
686    def visitKeyword(self, node):
687        self.emit('LOAD_CONST', node.name)
688        self.visit(node.expr)
689
690    def visitGlobal(self, node):
691        # no code to generate
692        pass
693
694    def visitName(self, node):
695        self.set_lineno(node)
696        self.loadName(node.name)
697
698    def visitPass(self, node):
699        self.set_lineno(node)
700
701    def visitImport(self, node):
702        self.set_lineno(node)
703        for name, alias in node.names:
704            if VERSION > 1:
705                self.emit('LOAD_CONST', None)
706            self.emit('IMPORT_NAME', name)
707            mod = string.split(name, ".")[0]
708            self.storeName(alias or mod)
709
710    def visitFrom(self, node):
711        self.set_lineno(node)
712        fromlist = map(lambda (name, alias): name, node.names)
713        if VERSION > 1:
714            self.emit('LOAD_CONST', tuple(fromlist))
715        self.emit('IMPORT_NAME', node.modname)
716        for name, alias in node.names:
717            if VERSION > 1:
718                if name == '*':
719                    self.namespace = 0
720                    self.emit('IMPORT_STAR')
721                    # There can only be one name w/ from ... import *
722                    assert len(node.names) == 1
723                    return
724                else:
725                    self.emit('IMPORT_FROM', name)
726                    self._resolveDots(name)
727                    self.storeName(alias or name)
728            else:
729                self.emit('IMPORT_FROM', name)
730        self.emit('POP_TOP')
731
732    def _resolveDots(self, name):
733        elts = string.split(name, ".")
734        if len(elts) == 1:
735            return
736        for elt in elts[1:]:
737            self.emit('LOAD_ATTR', elt)
738
739    def visitGetattr(self, node):
740        self.visit(node.expr)
741        self.emit('LOAD_ATTR', self.mangle(node.attrname))
742
743    # next five implement assignments
744
745    def visitAssign(self, node):
746        self.set_lineno(node)
747        self.visit(node.expr)
748        dups = len(node.nodes) - 1
749        for i in range(len(node.nodes)):
750            elt = node.nodes[i]
751            if i < dups:
752                self.emit('DUP_TOP')
753            if isinstance(elt, ast.Node):
754                self.visit(elt)
755
756    def visitAssName(self, node):
757        if node.flags == 'OP_ASSIGN':
758            self.storeName(node.name)
759        elif node.flags == 'OP_DELETE':
760            self.set_lineno(node)
761            self.delName(node.name)
762        else:
763            print "oops", node.flags
764
765    def visitAssAttr(self, node):
766        self.visit(node.expr)
767        if node.flags == 'OP_ASSIGN':
768            self.emit('STORE_ATTR', self.mangle(node.attrname))
769        elif node.flags == 'OP_DELETE':
770            self.emit('DELETE_ATTR', self.mangle(node.attrname))
771        else:
772            print "warning: unexpected flags:", node.flags
773            print node
774
775    def _visitAssSequence(self, node, op='UNPACK_SEQUENCE'):
776        if findOp(node) != 'OP_DELETE':
777            self.emit(op, len(node.nodes))
778        for child in node.nodes:
779            self.visit(child)
780
781    if VERSION > 1:
782        visitAssTuple = _visitAssSequence
783        visitAssList = _visitAssSequence
784    else:
785        def visitAssTuple(self, node):
786            self._visitAssSequence(node, 'UNPACK_TUPLE')
787
788        def visitAssList(self, node):
789            self._visitAssSequence(node, 'UNPACK_LIST')
790
791    # augmented assignment
792
793    def visitAugAssign(self, node):
794        self.set_lineno(node)
795        aug_node = wrap_aug(node.node)
796        self.visit(aug_node, "load")
797        self.visit(node.expr)
798        self.emit(self._augmented_opcode[node.op])
799        self.visit(aug_node, "store")
800
801    _augmented_opcode = {
802        '+=' : 'INPLACE_ADD',
803        '-=' : 'INPLACE_SUBTRACT',
804        '*=' : 'INPLACE_MULTIPLY',
805        '/=' : 'INPLACE_DIVIDE',
806        '//=': 'INPLACE_FLOOR_DIVIDE',
807        '%=' : 'INPLACE_MODULO',
808        '**=': 'INPLACE_POWER',
809        '>>=': 'INPLACE_RSHIFT',
810        '<<=': 'INPLACE_LSHIFT',
811        '&=' : 'INPLACE_AND',
812        '^=' : 'INPLACE_XOR',
813        '|=' : 'INPLACE_OR',
814        }
815
816    def visitAugName(self, node, mode):
817        if mode == "load":
818            self.loadName(node.name)
819        elif mode == "store":
820            self.storeName(node.name)
821
822    def visitAugGetattr(self, node, mode):
823        if mode == "load":
824            self.visit(node.expr)
825            self.emit('DUP_TOP')
826            self.emit('LOAD_ATTR', self.mangle(node.attrname))
827        elif mode == "store":
828            self.emit('ROT_TWO')
829            self.emit('STORE_ATTR', self.mangle(node.attrname))
830
831    def visitAugSlice(self, node, mode):
832        if mode == "load":
833            self.visitSlice(node, 1)
834        elif mode == "store":
835            slice = 0
836            if node.lower:
837                slice = slice | 1
838            if node.upper:
839                slice = slice | 2
840            if slice == 0:
841                self.emit('ROT_TWO')
842            elif slice == 3:
843                self.emit('ROT_FOUR')
844            else:
845                self.emit('ROT_THREE')
846            self.emit('STORE_SLICE+%d' % slice)
847
848    def visitAugSubscript(self, node, mode):
849        if len(node.subs) > 1:
850            raise SyntaxError, "augmented assignment to tuple is not possible"
851        if mode == "load":
852            self.visitSubscript(node, 1)
853        elif mode == "store":
854            self.emit('ROT_THREE')
855            self.emit('STORE_SUBSCR')
856
857    def visitExec(self, node):
858        self.visit(node.expr)
859        if node.locals is None:
860            self.emit('LOAD_CONST', None)
861        else:
862            self.visit(node.locals)
863        if node.globals is None:
864            self.emit('DUP_TOP')
865        else:
866            self.visit(node.globals)
867        self.emit('EXEC_STMT')
868
869    def visitCallFunc(self, node):
870        pos = 0
871        kw = 0
872        self.set_lineno(node)
873        self.visit(node.node)
874        for arg in node.args:
875            self.visit(arg)
876            if isinstance(arg, ast.Keyword):
877                kw = kw + 1
878            else:
879                pos = pos + 1
880        if node.star_args is not None:
881            self.visit(node.star_args)
882        if node.dstar_args is not None:
883            self.visit(node.dstar_args)
884        have_star = node.star_args is not None
885        have_dstar = node.dstar_args is not None
886        opcode = callfunc_opcode_info[have_star, have_dstar]
887        self.emit(opcode, kw << 8 | pos)
888
889    def visitPrint(self, node, newline=0):
890        self.set_lineno(node)
891        if node.dest:
892            self.visit(node.dest)
893        for child in node.nodes:
894            if node.dest:
895                self.emit('DUP_TOP')
896            self.visit(child)
897            if node.dest:
898                self.emit('ROT_TWO')
899                self.emit('PRINT_ITEM_TO')
900            else:
901                self.emit('PRINT_ITEM')
902        if node.dest and not newline:
903            self.emit('POP_TOP')
904
905    def visitPrintnl(self, node):
906        self.visitPrint(node, newline=1)
907        if node.dest:
908            self.emit('PRINT_NEWLINE_TO')
909        else:
910            self.emit('PRINT_NEWLINE')
911
912    def visitReturn(self, node):
913        self.set_lineno(node)
914        self.visit(node.value)
915        self.emit('RETURN_VALUE')
916
917    def visitYield(self, node):
918        self.set_lineno(node)
919        self.visit(node.value)
920        self.emit('YIELD_STMT')
921
922    # slice and subscript stuff
923
924    def visitSlice(self, node, aug_flag=None):
925        # aug_flag is used by visitAugSlice
926        self.visit(node.expr)
927        slice = 0
928        if node.lower:
929            self.visit(node.lower)
930            slice = slice | 1
931        if node.upper:
932            self.visit(node.upper)
933            slice = slice | 2
934        if aug_flag:
935            if slice == 0:
936                self.emit('DUP_TOP')
937            elif slice == 3:
938                self.emit('DUP_TOPX', 3)
939            else:
940                self.emit('DUP_TOPX', 2)
941        if node.flags == 'OP_APPLY':
942            self.emit('SLICE+%d' % slice)
943        elif node.flags == 'OP_ASSIGN':
944            self.emit('STORE_SLICE+%d' % slice)
945        elif node.flags == 'OP_DELETE':
946            self.emit('DELETE_SLICE+%d' % slice)
947        else:
948            print "weird slice", node.flags
949            raise
950
951    def visitSubscript(self, node, aug_flag=None):
952        self.visit(node.expr)
953        for sub in node.subs:
954            self.visit(sub)
955        if aug_flag:
956            self.emit('DUP_TOPX', 2)
957        if len(node.subs) > 1:
958            self.emit('BUILD_TUPLE', len(node.subs))
959        if node.flags == 'OP_APPLY':
960            self.emit('BINARY_SUBSCR')
961        elif node.flags == 'OP_ASSIGN':
962            self.emit('STORE_SUBSCR')
963        elif node.flags == 'OP_DELETE':
964            self.emit('DELETE_SUBSCR')
965
966    # binary ops
967
968    def binaryOp(self, node, op):
969        self.visit(node.left)
970        self.visit(node.right)
971        self.emit(op)
972
973    def visitAdd(self, node):
974        return self.binaryOp(node, 'BINARY_ADD')
975
976    def visitSub(self, node):
977        return self.binaryOp(node, 'BINARY_SUBTRACT')
978
979    def visitMul(self, node):
980        return self.binaryOp(node, 'BINARY_MULTIPLY')
981
982    def visitDiv(self, node):
983        return self.binaryOp(node, self._div_op)
984
985    def visitFloorDiv(self, node):
986        return self.binaryOp(node, 'BINARY_FLOOR_DIVIDE')
987
988    def visitMod(self, node):
989        return self.binaryOp(node, 'BINARY_MODULO')
990
991    def visitPower(self, node):
992        return self.binaryOp(node, 'BINARY_POWER')
993
994    def visitLeftShift(self, node):
995        return self.binaryOp(node, 'BINARY_LSHIFT')
996
997    def visitRightShift(self, node):
998        return self.binaryOp(node, 'BINARY_RSHIFT')
999
1000    # unary ops
1001
1002    def unaryOp(self, node, op):
1003        self.visit(node.expr)
1004        self.emit(op)
1005
1006    def visitInvert(self, node):
1007        return self.unaryOp(node, 'UNARY_INVERT')
1008
1009    def visitUnarySub(self, node):
1010        return self.unaryOp(node, 'UNARY_NEGATIVE')
1011
1012    def visitUnaryAdd(self, node):
1013        return self.unaryOp(node, 'UNARY_POSITIVE')
1014
1015    def visitUnaryInvert(self, node):
1016        return self.unaryOp(node, 'UNARY_INVERT')
1017
1018    def visitNot(self, node):
1019        return self.unaryOp(node, 'UNARY_NOT')
1020
1021    def visitBackquote(self, node):
1022        return self.unaryOp(node, 'UNARY_CONVERT')
1023
1024    # bit ops
1025
1026    def bitOp(self, nodes, op):
1027        self.visit(nodes[0])
1028        for node in nodes[1:]:
1029            self.visit(node)
1030            self.emit(op)
1031
1032    def visitBitand(self, node):
1033        return self.bitOp(node.nodes, 'BINARY_AND')
1034
1035    def visitBitor(self, node):
1036        return self.bitOp(node.nodes, 'BINARY_OR')
1037
1038    def visitBitxor(self, node):
1039        return self.bitOp(node.nodes, 'BINARY_XOR')
1040
1041    # object constructors
1042
1043    def visitEllipsis(self, node):
1044        self.emit('LOAD_CONST', Ellipsis)
1045
1046    def visitTuple(self, node):
1047        self.set_lineno(node)
1048        for elt in node.nodes:
1049            self.visit(elt)
1050        self.emit('BUILD_TUPLE', len(node.nodes))
1051
1052    def visitList(self, node):
1053        self.set_lineno(node)
1054        for elt in node.nodes:
1055            self.visit(elt)
1056        self.emit('BUILD_LIST', len(node.nodes))
1057
1058    def visitSliceobj(self, node):
1059        for child in node.nodes:
1060            self.visit(child)
1061        self.emit('BUILD_SLICE', len(node.nodes))
1062
1063    def visitDict(self, node):
1064        lineno = getattr(node, 'lineno', None)
1065        if lineno:
1066            self.emit('SET_LINENO', lineno)
1067        self.emit('BUILD_MAP', 0)
1068        for k, v in node.items:
1069            lineno2 = getattr(node, 'lineno', None)
1070            if lineno2 is not None and lineno != lineno2:
1071                self.emit('SET_LINENO', lineno2)
1072                lineno = lineno2
1073            self.emit('DUP_TOP')
1074            self.visit(v)
1075            self.emit('ROT_TWO')
1076            self.visit(k)
1077            self.emit('STORE_SUBSCR')
1078
1079class NestedScopeMixin:
1080    """Defines initClass() for nested scoping (Python 2.2-compatible)"""
1081    def initClass(self):
1082        self.__class__.NameFinder = LocalNameFinder
1083        self.__class__.FunctionGen = FunctionCodeGenerator
1084        self.__class__.ClassGen = ClassCodeGenerator
1085
1086class ModuleCodeGenerator(NestedScopeMixin, CodeGenerator):
1087    __super_init = CodeGenerator.__init__
1088
1089    scopes = None
1090
1091    def __init__(self, tree):
1092        self.graph = pyassem.PyFlowGraph("<module>", tree.filename)
1093        self.futures = future.find_futures(tree)
1094        self.__super_init()
1095        walk(tree, self)
1096
1097    def get_module(self):
1098        return self
1099
1100class AbstractFunctionCode:
1101    optimized = 1
1102    lambdaCount = 0
1103
1104    def __init__(self, func, scopes, isLambda, class_name, mod):
1105        self.class_name = class_name
1106        self.module = mod
1107        if isLambda:
1108            klass = FunctionCodeGenerator
1109            name = "<lambda.%d>" % klass.lambdaCount
1110            klass.lambdaCount = klass.lambdaCount + 1
1111        else:
1112            name = func.name
1113        args, hasTupleArg = generateArgList(func.argnames)
1114        self.graph = pyassem.PyFlowGraph(name, func.filename, args,
1115                                         optimized=1)
1116        self.isLambda = isLambda
1117        self.super_init()
1118
1119        if not isLambda and func.doc:
1120            self.setDocstring(func.doc)
1121
1122        lnf = walk(func.code, self.NameFinder(args), verbose=0)
1123        self.locals.push(lnf.getLocals())
1124        if func.varargs:
1125            self.graph.setFlag(CO_VARARGS)
1126        if func.kwargs:
1127            self.graph.setFlag(CO_VARKEYWORDS)
1128        self.set_lineno(func)
1129        if hasTupleArg:
1130            self.generateArgUnpack(func.argnames)
1131
1132    def get_module(self):
1133        return self.module
1134
1135    def finish(self):
1136        self.graph.startExitBlock()
1137        if not self.isLambda:
1138            self.emit('LOAD_CONST', None)
1139        self.emit('RETURN_VALUE')
1140
1141    def generateArgUnpack(self, args):
1142        for i in range(len(args)):
1143            arg = args[i]
1144            if type(arg) == types.TupleType:
1145                self.emit('LOAD_FAST', '.%d' % (i * 2))
1146                self.unpackSequence(arg)
1147
1148    def unpackSequence(self, tup):
1149        if VERSION > 1:
1150            self.emit('UNPACK_SEQUENCE', len(tup))
1151        else:
1152            self.emit('UNPACK_TUPLE', len(tup))
1153        for elt in tup:
1154            if type(elt) == types.TupleType:
1155                self.unpackSequence(elt)
1156            else:
1157                self._nameOp('STORE', elt)
1158
1159    unpackTuple = unpackSequence
1160
1161class FunctionCodeGenerator(NestedScopeMixin, AbstractFunctionCode,
1162                            CodeGenerator):
1163    super_init = CodeGenerator.__init__ # call be other init
1164    scopes = None
1165
1166    __super_init = AbstractFunctionCode.__init__
1167
1168    def __init__(self, func, scopes, isLambda, class_name, mod):
1169        self.scopes = scopes
1170        self.scope = scopes[func]
1171        self.__super_init(func, scopes, isLambda, class_name, mod)
1172        self.graph.setFreeVars(self.scope.get_free_vars())
1173        self.graph.setCellVars(self.scope.get_cell_vars())
1174        if self.graph.checkFlag(CO_GENERATOR_ALLOWED):
1175            if self.scope.generator is not None:
1176                self.graph.setFlag(CO_GENERATOR)
1177
1178class AbstractClassCode:
1179
1180    def __init__(self, klass, scopes, module):
1181        self.class_name = klass.name
1182        self.module = module
1183        self.graph = pyassem.PyFlowGraph(klass.name, klass.filename,
1184                                           optimized=0, klass=1)
1185        self.super_init()
1186        lnf = walk(klass.code, self.NameFinder(), verbose=0)
1187        self.locals.push(lnf.getLocals())
1188        self.graph.setFlag(CO_NEWLOCALS)
1189        if klass.doc:
1190            self.setDocstring(klass.doc)
1191
1192    def get_module(self):
1193        return self.module
1194
1195    def finish(self):
1196        self.graph.startExitBlock()
1197        self.emit('LOAD_LOCALS')
1198        self.emit('RETURN_VALUE')
1199
1200class ClassCodeGenerator(NestedScopeMixin, AbstractClassCode, CodeGenerator):
1201    super_init = CodeGenerator.__init__
1202    scopes = None
1203
1204    __super_init = AbstractClassCode.__init__
1205
1206    def __init__(self, klass, scopes, module):
1207        self.scopes = scopes
1208        self.scope = scopes[klass]
1209        self.__super_init(klass, scopes, module)
1210        self.graph.setFreeVars(self.scope.get_free_vars())
1211        self.graph.setCellVars(self.scope.get_cell_vars())
1212##        self.graph.setFlag(CO_NESTED)
1213
1214def generateArgList(arglist):
1215    """Generate an arg list marking TupleArgs"""
1216    args = []
1217    extra = []
1218    count = 0
1219    for i in range(len(arglist)):
1220        elt = arglist[i]
1221        if type(elt) == types.StringType:
1222            args.append(elt)
1223        elif type(elt) == types.TupleType:
1224            args.append(TupleArg(i * 2, elt))
1225            extra.extend(misc.flatten(elt))
1226            count = count + 1
1227        else:
1228            raise ValueError, "unexpect argument type:", elt
1229    return args + extra, count
1230
1231def findOp(node):
1232    """Find the op (DELETE, LOAD, STORE) in an AssTuple tree"""
1233    v = OpFinder()
1234    walk(node, v, verbose=0)
1235    return v.op
1236
1237class OpFinder:
1238    def __init__(self):
1239        self.op = None
1240    def visitAssName(self, node):
1241        if self.op is None:
1242            self.op = node.flags
1243        elif self.op != node.flags:
1244            raise ValueError, "mixed ops in stmt"
1245    visitAssAttr = visitAssName
1246
1247class Delegator:
1248    """Base class to support delegation for augmented assignment nodes
1249
1250    To generator code for augmented assignments, we use the following
1251    wrapper classes.  In visitAugAssign, the left-hand expression node
1252    is visited twice.  The first time the visit uses the normal method
1253    for that node .  The second time the visit uses a different method
1254    that generates the appropriate code to perform the assignment.
1255    These delegator classes wrap the original AST nodes in order to
1256    support the variant visit methods.
1257    """
1258    def __init__(self, obj):
1259        self.obj = obj
1260
1261    def __getattr__(self, attr):
1262        return getattr(self.obj, attr)
1263
1264class AugGetattr(Delegator):
1265    pass
1266
1267class AugName(Delegator):
1268    pass
1269
1270class AugSlice(Delegator):
1271    pass
1272
1273class AugSubscript(Delegator):
1274    pass
1275
1276wrapper = {
1277    ast.Getattr: AugGetattr,
1278    ast.Name: AugName,
1279    ast.Slice: AugSlice,
1280    ast.Subscript: AugSubscript,
1281    }
1282
1283def wrap_aug(node):
1284    return wrapper[node.__class__](node)
1285
1286if __name__ == "__main__":
1287    import sys
1288
1289    for file in sys.argv[1:]:
1290        compile(file)
1291