pycodegen.py revision 2051608616178ca6b73d5d59d887ddd67c5835e3 
1import os
2import marshal
3import stat
4import string
5import struct
6import types
7from cStringIO import StringIO
8
9from compiler import ast, parse, walk
10from compiler import pyassem, misc
11from compiler.pyassem import CO_VARARGS, CO_VARKEYWORDS, CO_NEWLOCALS, TupleArg
12
13callfunc_opcode_info = {
14    # (Have *args, Have **args) : opcode
15    (0,0) : "CALL_FUNCTION",
16    (1,0) : "CALL_FUNCTION_VAR",
17    (0,1) : "CALL_FUNCTION_KW",
18    (1,1) : "CALL_FUNCTION_VAR_KW",
19}
20
21def compile(filename):
22    f = open(filename)
23    buf = f.read()
24    f.close()
25    mod = Module(buf, filename)
26    mod.compile()
27    f = open(filename + "c", "wb")
28    mod.dump(f)
29    f.close()
30
31class Module:
32    def __init__(self, source, filename):
33        self.filename = filename
34        self.source = source
35        self.code = None
36
37    def compile(self):
38        ast = parse(self.source)
39        root, filename = os.path.split(self.filename)
40        gen = ModuleCodeGenerator(filename)
41        walk(ast, gen, 1)
42        self.code = gen.getCode()
43
44    def dump(self, f):
45        f.write(self.getPycHeader())
46        marshal.dump(self.code, f)
47
48    MAGIC = (50823 | (ord('\r')<<16) | (ord('\n')<<24))
49
50    def getPycHeader(self):
51        # compile.c uses marshal to write a long directly, with
52        # calling the interface that would also generate a 1-byte code
53        # to indicate the type of the value.  simplest way to get the
54        # same effect is to call marshal and then skip the code.
55        magic = marshal.dumps(self.MAGIC)[1:]
56        mtime = os.stat(self.filename)[stat.ST_MTIME]
57        mtime = struct.pack('i', mtime)
58        return magic + mtime
59
60class CodeGenerator:
61
62    optimized = 0 # is namespace access optimized?
63
64    def __init__(self, filename):
65## Subclasses must define a constructor that intializes self.graph
66## before calling this init function
67##         self.graph = pyassem.PyFlowGraph()
68        self.filename = filename
69        self.locals = misc.Stack()
70        self.loops = misc.Stack()
71        self.curStack = 0
72        self.maxStack = 0
73        self._setupGraphDelegation()
74
75    def _setupGraphDelegation(self):
76        self.emit = self.graph.emit
77        self.newBlock = self.graph.newBlock
78        self.startBlock = self.graph.startBlock
79        self.nextBlock = self.graph.nextBlock
80        self.setDocstring = self.graph.setDocstring
81
82    def getCode(self):
83        """Return a code object"""
84        return self.graph.getCode()
85
86    # Next five methods handle name access
87
88    def isLocalName(self, name):
89        return self.locals.top().has_elt(name)
90
91    def storeName(self, name):
92        self._nameOp('STORE', name)
93
94    def loadName(self, name):
95        self._nameOp('LOAD', name)
96
97    def delName(self, name):
98        self._nameOp('DELETE', name)
99
100    def _nameOp(self, prefix, name):
101        if not self.optimized:
102            self.emit(prefix + '_NAME', name)
103            return
104        if self.isLocalName(name):
105            self.emit(prefix + '_FAST', name)
106        else:
107            self.emit(prefix + '_GLOBAL', name)
108
109    def set_lineno(self, node):
110        """Emit SET_LINENO if node has lineno attribute
111
112        Returns true if SET_LINENO was emitted.
113
114        There are no rules for when an AST node should have a lineno
115        attribute.  The transformer and AST code need to be reviewed
116        and a consistent policy implemented and documented.  Until
117        then, this method works around missing line numbers.
118        """
119        lineno = getattr(node, 'lineno', None)
120        if lineno is not None:
121            self.emit('SET_LINENO', lineno)
122            return 1
123        return 0
124
125    # The first few visitor methods handle nodes that generator new
126    # code objects
127
128    def visitModule(self, node):
129        lnf = walk(node.node, LocalNameFinder(), 0)
130        self.locals.push(lnf.getLocals())
131        self.setDocstring(node.doc)
132        self.visit(node.node)
133        self.emit('LOAD_CONST', None)
134        self.emit('RETURN_VALUE')
135
136    def visitFunction(self, node):
137        self._visitFuncOrLambda(node, isLambda=0)
138        self.storeName(node.name)
139
140    def visitLambda(self, node):
141        self._visitFuncOrLambda(node, isLambda=1)
142##        self.storeName("<lambda>")
143
144    def _visitFuncOrLambda(self, node, isLambda):
145        gen = FunctionCodeGenerator(node, self.filename, isLambda)
146        walk(node.code, gen)
147        gen.finish()
148        self.set_lineno(node)
149        for default in node.defaults:
150            self.visit(default)
151        self.emit('LOAD_CONST', gen.getCode())
152        self.emit('MAKE_FUNCTION', len(node.defaults))
153
154    def visitClass(self, node):
155        gen = ClassCodeGenerator(node, self.filename)
156        walk(node.code, gen)
157        gen.finish()
158        self.set_lineno(node)
159        self.emit('LOAD_CONST', node.name)
160        for base in node.bases:
161            self.visit(base)
162        self.emit('BUILD_TUPLE', len(node.bases))
163        self.emit('LOAD_CONST', gen.getCode())
164        self.emit('MAKE_FUNCTION', 0)
165        self.emit('CALL_FUNCTION', 0)
166        self.emit('BUILD_CLASS')
167        self.storeName(node.name)
168
169    # The rest are standard visitor methods
170
171    # The next few implement control-flow statements
172
173    def visitIf(self, node):
174        end = self.newBlock()
175        numtests = len(node.tests)
176        for i in range(numtests):
177            test, suite = node.tests[i]
178            self.set_lineno(test)
179            self.visit(test)
180##            if i == numtests - 1 and not node.else_:
181##                nextTest = end
182##            else:
183##                nextTest = self.newBlock()
184            nextTest = self.newBlock()
185            self.emit('JUMP_IF_FALSE', nextTest)
186            self.nextBlock()
187            self.emit('POP_TOP')
188            self.visit(suite)
189            self.emit('JUMP_FORWARD', end)
190            self.nextBlock(nextTest)
191            self.emit('POP_TOP')
192        if node.else_:
193            self.visit(node.else_)
194        self.nextBlock(end)
195
196    def visitWhile(self, node):
197        self.set_lineno(node)
198
199        loop = self.newBlock()
200        else_ = self.newBlock()
201
202        after = self.newBlock()
203        self.emit('SETUP_LOOP', after)
204
205        self.nextBlock(loop)
206        self.loops.push(loop)
207
208        self.set_lineno(node)
209        self.visit(node.test)
210        self.emit('JUMP_IF_FALSE', else_ or after)
211
212        self.nextBlock()
213        self.emit('POP_TOP')
214        self.visit(node.body)
215        self.emit('JUMP_ABSOLUTE', loop)
216
217        self.startBlock(else_) # or just the POPs if not else clause
218        self.emit('POP_TOP')
219        self.emit('POP_BLOCK')
220        if node.else_:
221            self.visit(node.else_)
222        self.loops.pop()
223        self.nextBlock(after)
224
225    def visitFor(self, node):
226        start = self.newBlock()
227        anchor = self.newBlock()
228        after = self.newBlock()
229        self.loops.push(start)
230
231        self.set_lineno(node)
232        self.emit('SETUP_LOOP', after)
233        self.visit(node.list)
234        self.visit(ast.Const(0))
235        self.nextBlock(start)
236        self.set_lineno(node)
237        self.emit('FOR_LOOP', anchor)
238        self.visit(node.assign)
239        self.visit(node.body)
240        self.emit('JUMP_ABSOLUTE', start)
241        self.nextBlock(anchor)
242        self.emit('POP_BLOCK')
243        if node.else_:
244            self.visit(node.else_)
245        self.loops.pop()
246        self.nextBlock(after)
247
248    def visitBreak(self, node):
249        if not self.loops:
250            raise SyntaxError, "'break' outside loop (%s, %d)" % \
251                  (self.filename, node.lineno)
252        self.set_lineno(node)
253        self.emit('BREAK_LOOP')
254
255    def visitContinue(self, node):
256        if not self.loops:
257            raise SyntaxError, "'continue' outside loop (%s, %d)" % \
258                  (self.filename, node.lineno)
259        l = self.loops.top()
260        self.set_lineno(node)
261        self.emit('JUMP_ABSOLUTE', l)
262        self.nextBlock()
263
264    def visitTest(self, node, jump):
265        end = self.newBlock()
266        for child in node.nodes[:-1]:
267            self.visit(child)
268            self.emit(jump, end)
269            self.nextBlock()
270            self.emit('POP_TOP')
271        self.visit(node.nodes[-1])
272        self.nextBlock(end)
273
274    def visitAnd(self, node):
275        self.visitTest(node, 'JUMP_IF_FALSE')
276
277    def visitOr(self, node):
278        self.visitTest(node, 'JUMP_IF_TRUE')
279
280    def visitCompare(self, node):
281        self.visit(node.expr)
282        cleanup = self.newBlock()
283        for op, code in node.ops[:-1]:
284            self.visit(code)
285            self.emit('DUP_TOP')
286            self.emit('ROT_THREE')
287            self.emit('COMPARE_OP', op)
288            self.emit('JUMP_IF_FALSE', cleanup)
289            self.nextBlock()
290            self.emit('POP_TOP')
291        # now do the last comparison
292        if node.ops:
293            op, code = node.ops[-1]
294            self.visit(code)
295            self.emit('COMPARE_OP', op)
296        if len(node.ops) > 1:
297            end = self.newBlock()
298            self.emit('JUMP_FORWARD', end)
299            self.nextBlock(cleanup)
300            self.emit('ROT_TWO')
301            self.emit('POP_TOP')
302            self.nextBlock(end)
303
304    # exception related
305
306    def visitAssert(self, node):
307        # XXX would be interesting to implement this via a
308        # transformation of the AST before this stage
309        end = self.newBlock()
310        self.set_lineno(node)
311        # XXX __debug__ and AssertionError appear to be special cases
312        # -- they are always loaded as globals even if there are local
313        # names.  I guess this is a sort of renaming op.
314        self.emit('LOAD_GLOBAL', '__debug__')
315        self.emit('JUMP_IF_FALSE', end)
316        self.nextBlock()
317        self.emit('POP_TOP')
318        self.visit(node.test)
319        self.emit('JUMP_IF_TRUE', end)
320        self.nextBlock()
321        self.emit('LOAD_GLOBAL', 'AssertionError')
322        self.visit(node.fail)
323        self.emit('RAISE_VARARGS', 2)
324        self.nextBlock(end)
325        self.emit('POP_TOP')
326
327    def visitRaise(self, node):
328        self.set_lineno(node)
329        n = 0
330        if node.expr1:
331            self.visit(node.expr1)
332            n = n + 1
333        if node.expr2:
334            self.visit(node.expr2)
335            n = n + 1
336        if node.expr3:
337            self.visit(node.expr3)
338            n = n + 1
339        self.emit('RAISE_VARARGS', n)
340
341    def visitTryExcept(self, node):
342        handlers = self.newBlock()
343        end = self.newBlock()
344        if node.else_:
345            lElse = self.newBlock()
346        else:
347            lElse = end
348        self.set_lineno(node)
349        self.emit('SETUP_EXCEPT', handlers)
350        self.visit(node.body)
351        self.emit('POP_BLOCK')
352        self.emit('JUMP_FORWARD', lElse)
353        self.nextBlock(handlers)
354
355        last = len(node.handlers) - 1
356        for i in range(len(node.handlers)):
357            expr, target, body = node.handlers[i]
358            self.set_lineno(expr)
359            if expr:
360                self.emit('DUP_TOP')
361                self.visit(expr)
362                self.emit('COMPARE_OP', 'exception match')
363                next = self.newBlock()
364                self.emit('JUMP_IF_FALSE', next)
365                self.nextBlock()
366                self.emit('POP_TOP')
367            self.emit('POP_TOP')
368            if target:
369                self.visit(target)
370            else:
371                self.emit('POP_TOP')
372            self.emit('POP_TOP')
373            self.visit(body)
374            self.emit('JUMP_FORWARD', end)
375            if expr:
376                self.nextBlock(next)
377            self.emit('POP_TOP')
378        self.emit('END_FINALLY')
379        if node.else_:
380            self.nextBlock(lElse)
381            self.visit(node.else_)
382        self.nextBlock(end)
383
384    def visitTryFinally(self, node):
385        final = self.newBlock()
386        self.set_lineno(node)
387        self.emit('SETUP_FINALLY', final)
388        self.visit(node.body)
389        self.emit('POP_BLOCK')
390        self.emit('LOAD_CONST', None)
391        self.nextBlock(final)
392        self.visit(node.final)
393        self.emit('END_FINALLY')
394
395    # misc
396
397##     def visitStmt(self, node):
398##         # nothing to do except walk the children
399##         pass
400
401    def visitDiscard(self, node):
402        self.visit(node.expr)
403        self.emit('POP_TOP')
404
405    def visitConst(self, node):
406        self.emit('LOAD_CONST', node.value)
407
408    def visitKeyword(self, node):
409        self.emit('LOAD_CONST', node.name)
410        self.visit(node.expr)
411
412    def visitGlobal(self, node):
413        # no code to generate
414        pass
415
416    def visitName(self, node):
417        self.loadName(node.name)
418
419    def visitPass(self, node):
420        self.set_lineno(node)
421
422    def visitImport(self, node):
423        self.set_lineno(node)
424        for name, alias in node.names:
425            self.emit('LOAD_CONST', None)
426            self.emit('IMPORT_NAME', name)
427            self._resolveDots(name)
428            self.storeName(alias or name)
429
430    def visitFrom(self, node):
431        self.set_lineno(node)
432        fromlist = map(lambda (name, alias): name, node.names)
433        self.emit('LOAD_CONST', tuple(fromlist))
434        self.emit('IMPORT_NAME', node.modname)
435        for name, alias in node.names:
436            if name == '*':
437                self.namespace = 0
438            self.emit('IMPORT_FROM', name)
439            self._resolveDots(name)
440            self.storeName(alias or name)
441        self.emit('POP_TOP')
442
443    def _resolveDots(self, name):
444        elts = string.split(name, ".")
445        if len(elts) == 1:
446            return
447        for elt in elts[1:]:
448            self.emit('LOAD_ATTR', elt)
449
450    def visitGetattr(self, node):
451        self.visit(node.expr)
452        self.emit('LOAD_ATTR', node.attrname)
453
454    # next five implement assignments
455
456    def visitAssign(self, node):
457        self.set_lineno(node)
458        self.visit(node.expr)
459        dups = len(node.nodes) - 1
460        for i in range(len(node.nodes)):
461            elt = node.nodes[i]
462            if i < dups:
463                self.emit('DUP_TOP')
464            if isinstance(elt, ast.Node):
465                self.visit(elt)
466
467    def visitAssName(self, node):
468        if node.flags == 'OP_ASSIGN':
469            self.storeName(node.name)
470        elif node.flags == 'OP_DELETE':
471            self.delName(node.name)
472        else:
473            print "oops", node.flags
474
475    def visitAssAttr(self, node):
476        self.visit(node.expr)
477        if node.flags == 'OP_ASSIGN':
478            self.emit('STORE_ATTR', node.attrname)
479        elif node.flags == 'OP_DELETE':
480            self.emit('DELETE_ATTR', node.attrname)
481        else:
482            print "warning: unexpected flags:", node.flags
483            print node
484
485    def visitAssTuple(self, node):
486        if findOp(node) != 'OP_DELETE':
487            self.emit('UNPACK_SEQUENCE', len(node.nodes))
488        for child in node.nodes:
489            self.visit(child)
490
491    visitAssList = visitAssTuple
492
493    def visitExec(self, node):
494        self.visit(node.expr)
495        if node.locals is None:
496            self.emit('LOAD_CONST', None)
497        else:
498            self.visit(node.locals)
499        if node.globals is None:
500            self.emit('DUP_TOP')
501        else:
502            self.visit(node.globals)
503        self.emit('EXEC_STMT')
504
505    def visitCallFunc(self, node):
506        pos = 0
507        kw = 0
508        self.set_lineno(node)
509        self.visit(node.node)
510        for arg in node.args:
511            self.visit(arg)
512            if isinstance(arg, ast.Keyword):
513                kw = kw + 1
514            else:
515                pos = pos + 1
516        if node.star_args is not None:
517            self.visit(node.star_args)
518        if node.dstar_args is not None:
519            self.visit(node.dstar_args)
520        have_star = node.star_args is not None
521        have_dstar = node.dstar_args is not None
522        opcode = callfunc_opcode_info[have_star, have_dstar]
523        self.emit(opcode, kw << 8 | pos)
524
525    def visitPrint(self, node):
526        self.set_lineno(node)
527        for child in node.nodes:
528            self.visit(child)
529            self.emit('PRINT_ITEM')
530
531    def visitPrintnl(self, node):
532        self.visitPrint(node)
533        self.emit('PRINT_NEWLINE')
534
535    def visitReturn(self, node):
536        self.set_lineno(node)
537        self.visit(node.value)
538        self.emit('RETURN_VALUE')
539
540    # slice and subscript stuff
541
542    def visitSlice(self, node):
543        self.visit(node.expr)
544        slice = 0
545        if node.lower:
546            self.visit(node.lower)
547            slice = slice | 1
548        if node.upper:
549            self.visit(node.upper)
550            slice = slice | 2
551        if node.flags == 'OP_APPLY':
552            self.emit('SLICE+%d' % slice)
553        elif node.flags == 'OP_ASSIGN':
554            self.emit('STORE_SLICE+%d' % slice)
555        elif node.flags == 'OP_DELETE':
556            self.emit('DELETE_SLICE+%d' % slice)
557        else:
558            print "weird slice", node.flags
559            raise
560
561    def visitSubscript(self, node):
562        self.visit(node.expr)
563        for sub in node.subs:
564            self.visit(sub)
565        if len(node.subs) > 1:
566            self.emit('BUILD_TUPLE', len(node.subs))
567        if node.flags == 'OP_APPLY':
568            self.emit('BINARY_SUBSCR')
569        elif node.flags == 'OP_ASSIGN':
570            self.emit('STORE_SUBSCR')
571        elif node.flags == 'OP_DELETE':
572            self.emit('DELETE_SUBSCR')
573
574    # binary ops
575
576    def binaryOp(self, node, op):
577        self.visit(node.left)
578        self.visit(node.right)
579        self.emit(op)
580
581    def visitAdd(self, node):
582        return self.binaryOp(node, 'BINARY_ADD')
583
584    def visitSub(self, node):
585        return self.binaryOp(node, 'BINARY_SUBTRACT')
586
587    def visitMul(self, node):
588        return self.binaryOp(node, 'BINARY_MULTIPLY')
589
590    def visitDiv(self, node):
591        return self.binaryOp(node, 'BINARY_DIVIDE')
592
593    def visitMod(self, node):
594        return self.binaryOp(node, 'BINARY_MODULO')
595
596    def visitPower(self, node):
597        return self.binaryOp(node, 'BINARY_POWER')
598
599    def visitLeftShift(self, node):
600        return self.binaryOp(node, 'BINARY_LSHIFT')
601
602    def visitRightShift(self, node):
603        return self.binaryOp(node, 'BINARY_RSHIFT')
604
605    # unary ops
606
607    def unaryOp(self, node, op):
608        self.visit(node.expr)
609        self.emit(op)
610
611    def visitInvert(self, node):
612        return self.unaryOp(node, 'UNARY_INVERT')
613
614    def visitUnarySub(self, node):
615        return self.unaryOp(node, 'UNARY_NEGATIVE')
616
617    def visitUnaryAdd(self, node):
618        return self.unaryOp(node, 'UNARY_POSITIVE')
619
620    def visitUnaryInvert(self, node):
621        return self.unaryOp(node, 'UNARY_INVERT')
622
623    def visitNot(self, node):
624        return self.unaryOp(node, 'UNARY_NOT')
625
626    def visitBackquote(self, node):
627        return self.unaryOp(node, 'UNARY_CONVERT')
628
629    # bit ops
630
631    def bitOp(self, nodes, op):
632        self.visit(nodes[0])
633        for node in nodes[1:]:
634            self.visit(node)
635            self.emit(op)
636
637    def visitBitand(self, node):
638        return self.bitOp(node.nodes, 'BINARY_AND')
639
640    def visitBitor(self, node):
641        return self.bitOp(node.nodes, 'BINARY_OR')
642
643    def visitBitxor(self, node):
644        return self.bitOp(node.nodes, 'BINARY_XOR')
645
646    # object constructors
647
648    def visitEllipsis(self, node):
649        self.emit('LOAD_CONST', Ellipsis)
650
651    def visitTuple(self, node):
652        for elt in node.nodes:
653            self.visit(elt)
654        self.emit('BUILD_TUPLE', len(node.nodes))
655
656    def visitList(self, node):
657        for elt in node.nodes:
658            self.visit(elt)
659        self.emit('BUILD_LIST', len(node.nodes))
660
661    def visitSliceobj(self, node):
662        for child in node.nodes:
663            self.visit(child)
664        self.emit('BUILD_SLICE', len(node.nodes))
665
666    def visitDict(self, node):
667        lineno = getattr(node, 'lineno', None)
668        if lineno:
669            set.emit('SET_LINENO', lineno)
670        self.emit('BUILD_MAP', 0)
671        for k, v in node.items:
672            lineno2 = getattr(node, 'lineno', None)
673            if lineno2 is not None and lineno != lineno2:
674                self.emit('SET_LINENO', lineno2)
675                lineno = lineno2
676            self.emit('DUP_TOP')
677            self.visit(v)
678            self.emit('ROT_TWO')
679            self.visit(k)
680            self.emit('STORE_SUBSCR')
681
682class ModuleCodeGenerator(CodeGenerator):
683    super_init = CodeGenerator.__init__
684
685    def __init__(self, filename):
686        # XXX <module> is ? in compile.c
687        self.graph = pyassem.PyFlowGraph("<module>", filename)
688        self.super_init(filename)
689
690class FunctionCodeGenerator(CodeGenerator):
691    super_init = CodeGenerator.__init__
692
693    optimized = 1
694    lambdaCount = 0
695
696    def __init__(self, func, filename, isLambda=0):
697        if isLambda:
698            klass = FunctionCodeGenerator
699            name = "<lambda.%d>" % klass.lambdaCount
700            klass.lambdaCount = klass.lambdaCount + 1
701        else:
702            name = func.name
703        args, hasTupleArg = generateArgList(func.argnames)
704        self.graph = pyassem.PyFlowGraph(name, filename, args,
705                                           optimized=1)
706        self.isLambda = isLambda
707        self.super_init(filename)
708
709        lnf = walk(func.code, LocalNameFinder(args), 0)
710        self.locals.push(lnf.getLocals())
711        if func.varargs:
712            self.graph.setFlag(CO_VARARGS)
713        if func.kwargs:
714            self.graph.setFlag(CO_VARKEYWORDS)
715        self.set_lineno(func)
716        if hasTupleArg:
717            self.generateArgUnpack(func.argnames)
718
719    def finish(self):
720        self.graph.startExitBlock()
721        if not self.isLambda:
722            self.emit('LOAD_CONST', None)
723        self.emit('RETURN_VALUE')
724
725    def generateArgUnpack(self, args):
726        count = 0
727        for arg in args:
728            if type(arg) == types.TupleType:
729                self.emit('LOAD_FAST', '.nested%d' % count)
730                count = count + 1
731                self.unpackSequence(arg)
732
733    def unpackSequence(self, tup):
734        self.emit('UNPACK_SEQUENCE', len(tup))
735        for elt in tup:
736            if type(elt) == types.TupleType:
737                self.unpackSequence(elt)
738            else:
739                self.emit('STORE_FAST', elt)
740
741    unpackTuple = unpackSequence
742
743class ClassCodeGenerator(CodeGenerator):
744    super_init = CodeGenerator.__init__
745
746    def __init__(self, klass, filename):
747        self.graph = pyassem.PyFlowGraph(klass.name, filename,
748                                           optimized=0)
749        self.super_init(filename)
750        lnf = walk(klass.code, LocalNameFinder(), 0)
751        self.locals.push(lnf.getLocals())
752        self.graph.setFlag(CO_NEWLOCALS)
753
754    def finish(self):
755        self.graph.startExitBlock()
756        self.emit('LOAD_LOCALS')
757        self.emit('RETURN_VALUE')
758
759
760def generateArgList(arglist):
761    """Generate an arg list marking TupleArgs"""
762    args = []
763    extra = []
764    count = 0
765    for elt in arglist:
766        if type(elt) == types.StringType:
767            args.append(elt)
768        elif type(elt) == types.TupleType:
769            args.append(TupleArg(count, elt))
770            count = count + 1
771            extra.extend(misc.flatten(elt))
772        else:
773            raise ValueError, "unexpect argument type:", elt
774    return args + extra, count
775
776class LocalNameFinder:
777    """Find local names in scope"""
778    def __init__(self, names=()):
779        self.names = misc.Set()
780        self.globals = misc.Set()
781        for name in names:
782            self.names.add(name)
783
784    def getLocals(self):
785        for elt in self.globals.elements():
786            if self.names.has_elt(elt):
787                self.names.remove(elt)
788        return self.names
789
790    def visitDict(self, node):
791        pass
792
793    def visitGlobal(self, node):
794        for name in node.names:
795            self.globals.add(name)
796
797    def visitFunction(self, node):
798        self.names.add(node.name)
799
800    def visitLambda(self, node):
801        pass
802
803    def visitImport(self, node):
804        for name, alias in node.names:
805            self.names.add(alias or name)
806
807    def visitFrom(self, node):
808        for name, alias in node.names:
809            self.names.add(alias or name)
810
811    def visitClass(self, node):
812        self.names.add(node.name)
813
814    def visitAssName(self, node):
815        self.names.add(node.name)
816
817def findOp(node):
818    """Find the op (DELETE, LOAD, STORE) in an AssTuple tree"""
819    v = OpFinder()
820    walk(node, v, 0)
821    return v.op
822
823class OpFinder:
824    def __init__(self):
825        self.op = None
826    def visitAssName(self, node):
827        if self.op is None:
828            self.op = node.flags
829        elif self.op != node.flags:
830            raise ValueError, "mixed ops in stmt"
831
832if __name__ == "__main__":
833    import sys
834
835    for file in sys.argv[1:]:
836        compile(file)
837