Lib/compiler/pycodegen.py

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