Lib/compiler/pycodegen.py

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