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