Lib/compiler/pycodegen.py

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