import imp

import os

import marshal

import struct

import sys

from cStringIO import StringIO

from compiler import ast, parse, walk, syntax

from compiler import pyassem, misc, future, symbols

from compiler.consts import SC_LOCAL, SC_GLOBAL, SC_FREE, SC_CELL

from compiler.consts import (CO_VARARGS, CO_VARKEYWORDS, CO_NEWLOCALS,

     CO_NESTED, CO_GENERATOR, CO_FUTURE_DIVISION,

     CO_FUTURE_ABSIMPORT, CO_FUTURE_WITH_STATEMENT, CO_FUTURE_PRINT_FUNCTION)

from compiler.pyassem import TupleArg

# XXX The version-specific code can go, since this code only works with 2.x.

# Do we have Python 1.x or Python 2.x?

try:

    VERSION = sys.version_info[0]

except AttributeError:

    VERSION = 1

callfunc_opcode_info = {

    # (Have *args, Have **args) : opcode

    (0,0) : "CALL_FUNCTION",

    (1,0) : "CALL_FUNCTION_VAR",

    (0,1) : "CALL_FUNCTION_KW",

    (1,1) : "CALL_FUNCTION_VAR_KW",

}

LOOP = 1

EXCEPT = 2

TRY_FINALLY = 3

END_FINALLY = 4

def compileFile(filename, display=0):

    f = open(filename, 'U')

    buf = f.read()

    f.close()

    mod = Module(buf, filename)

    try:

        mod.compile(display)

    except SyntaxError:

        raise

    else:

        f = open(filename + "c", "wb")

        mod.dump(f)

        f.close()

def compile(source, filename, mode, flags=None, dont_inherit=None):

    """Replacement for builtin compile() function"""

    if flags is not None or dont_inherit is not None:

        raise RuntimeError, "not implemented yet"

    if mode == "single":

        gen = Interactive(source, filename)

    elif mode == "exec":

        gen = Module(source, filename)

    elif mode == "eval":

        gen = Expression(source, filename)

    else:

        raise ValueError("compile() 3rd arg must be 'exec' or "

                         "'eval' or 'single'")

    gen.compile()

    return gen.code

class AbstractCompileMode:

    mode = None # defined by subclass

    def __init__(self, source, filename):

        self.source = source

        self.filename = filename

        self.code = None

    def _get_tree(self):

        tree = parse(self.source, self.mode)

        misc.set_filename(self.filename, tree)

        syntax.check(tree)

        return tree

    def compile(self):

        pass # implemented by subclass

    def getCode(self):

        return self.code

class Expression(AbstractCompileMode):

    mode = "eval"

    def compile(self):

        tree = self._get_tree()

        gen = ExpressionCodeGenerator(tree)

        self.code = gen.getCode()

class Interactive(AbstractCompileMode):

    mode = "single"

    def compile(self):

        tree = self._get_tree()

        gen = InteractiveCodeGenerator(tree)

        self.code = gen.getCode()

class Module(AbstractCompileMode):

    mode = "exec"

    def compile(self, display=0):

        tree = self._get_tree()

        gen = ModuleCodeGenerator(tree)

        if display:

            import pprint

            print pprint.pprint(tree)

        self.code = gen.getCode()

    def dump(self, f):

        f.write(self.getPycHeader())

        marshal.dump(self.code, f)

    MAGIC = imp.get_magic()

    def getPycHeader(self):

        # compile.c uses marshal to write a long directly, with

        # calling the interface that would also generate a 1-byte code

        # to indicate the type of the value.  simplest way to get the

        # same effect is to call marshal and then skip the code.

        mtime = os.path.getmtime(self.filename)

        mtime = struct.pack('<i', mtime)

        return self.MAGIC + mtime

class LocalNameFinder:

    """Find local names in scope"""

    def __init__(self, names=()):

        self.names = misc.Set()

        self.globals = misc.Set()

        for name in names:

            self.names.add(name)

    # XXX list comprehensions and for loops

    def getLocals(self):

        for elt in self.globals.elements():

            if self.names.has_elt(elt):

                self.names.remove(elt)

        return self.names

    def visitDict(self, node):

        pass

    def visitGlobal(self, node):

        for name in node.names:

            self.globals.add(name)

    def visitFunction(self, node):

        self.names.add(node.name)

    def visitLambda(self, node):

        pass

    def visitImport(self, node):

        for name, alias in node.names:

            self.names.add(alias or name)

    def visitFrom(self, node):

        for name, alias in node.names:

            self.names.add(alias or name)

    def visitClass(self, node):

        self.names.add(node.name)

    def visitAssName(self, node):

        self.names.add(node.name)

def is_constant_false(node):

    if isinstance(node, ast.Const):

        if not node.value:

            return 1

    return 0

class CodeGenerator:

    """Defines basic code generator for Python bytecode

    This class is an abstract base class.  Concrete subclasses must

    define an __init__() that defines self.graph and then calls the

    __init__() defined in this class.

    The concrete class must also define the class attributes

    NameFinder, FunctionGen, and ClassGen.  These attributes can be

    defined in the initClass() method, which is a hook for

    initializing these methods after all the classes have been

    defined.

    """

    optimized = 0 # is namespace access optimized?

    __initialized = None

    class_name = None # provide default for instance variable

    def __init__(self):

        if self.__initialized is None:

            self.initClass()

            self.__class__.__initialized = 1

        self.checkClass()

        self.locals = misc.Stack()

        self.setups = misc.Stack()

        self.last_lineno = None

        self._setupGraphDelegation()

        self._div_op = "BINARY_DIVIDE"

        # XXX set flags based on future features

        futures = self.get_module().futures

        for feature in futures:

            if feature == "division":

                self.graph.setFlag(CO_FUTURE_DIVISION)

                self._div_op = "BINARY_TRUE_DIVIDE"

            elif feature == "absolute_import":

                self.graph.setFlag(CO_FUTURE_ABSIMPORT)

            elif feature == "with_statement":

                self.graph.setFlag(CO_FUTURE_WITH_STATEMENT)

            elif feature == "print_function":

                self.graph.setFlag(CO_FUTURE_PRINT_FUNCTION)

    def initClass(self):

        """This method is called once for each class"""

    def checkClass(self):

        """Verify that class is constructed correctly"""

        try:

            assert hasattr(self, 'graph')

            assert getattr(self, 'NameFinder')

            assert getattr(self, 'FunctionGen')

            assert getattr(self, 'ClassGen')

        except AssertionError, msg:

            intro = "Bad class construction for %s" % self.__class__.__name__

            raise AssertionError, intro

    def _setupGraphDelegation(self):

        self.emit = self.graph.emit

        self.newBlock = self.graph.newBlock

        self.startBlock = self.graph.startBlock

        self.nextBlock = self.graph.nextBlock

        self.setDocstring = self.graph.setDocstring

    def getCode(self):

        """Return a code object"""

        return self.graph.getCode()

    def mangle(self, name):

        if self.class_name is not None:

            return misc.mangle(name, self.class_name)

        else:

            return name

    def parseSymbols(self, tree):

        s = symbols.SymbolVisitor()

        walk(tree, s)

        return s.scopes

    def get_module(self):

        raise RuntimeError, "should be implemented by subclasses"

    # Next five methods handle name access

    def isLocalName(self, name):

        return self.locals.top().has_elt(name)

    def storeName(self, name):

        self._nameOp('STORE', name)

    def loadName(self, name):

        self._nameOp('LOAD', name)

    def delName(self, name):

        self._nameOp('DELETE', name)

    def _nameOp(self, prefix, name):

        name = self.mangle(name)

        scope = self.scope.check_name(name)

        if scope == SC_LOCAL:

            if not self.optimized:

                self.emit(prefix + '_NAME', name)

            else:

                self.emit(prefix + '_FAST', name)

        elif scope == SC_GLOBAL:

            if not self.optimized:

                self.emit(prefix + '_NAME', name)

            else:

                self.emit(prefix + '_GLOBAL', name)

        elif scope == SC_FREE or scope == SC_CELL:

            self.emit(prefix + '_DEREF', name)

        else:

            raise RuntimeError, "unsupported scope for var %s: %d" % \

                  (name, scope)

    def _implicitNameOp(self, prefix, name):

        """Emit name ops for names generated implicitly by for loops

        The interpreter generates names that start with a period or

        dollar sign.  The symbol table ignores these names because

        they aren't present in the program text.

        """

        if self.optimized:

            self.emit(prefix + '_FAST', name)

        else:

            self.emit(prefix + '_NAME', name)

    # The set_lineno() function and the explicit emit() calls for

    # SET_LINENO below are only used to generate the line number table.

    # As of Python 2.3, the interpreter does not have a SET_LINENO

    # instruction.  pyassem treats SET_LINENO opcodes as a special case.

    def set_lineno(self, node, force=False):

        """Emit SET_LINENO if necessary.

        The instruction is considered necessary if the node has a

        lineno attribute and it is different than the last lineno

        emitted.

        Returns true if SET_LINENO was emitted.

        There are no rules for when an AST node should have a lineno

        attribute.  The transformer and AST code need to be reviewed

        and a consistent policy implemented and documented.  Until

        then, this method works around missing line numbers.

        """

        lineno = getattr(node, 'lineno', None)

        if lineno is not None and (lineno != self.last_lineno

                                   or force):

            self.emit('SET_LINENO', lineno)

            self.last_lineno = lineno

            return True

        return False

    # The first few visitor methods handle nodes that generator new

    # code objects.  They use class attributes to determine what

    # specialized code generators to use.

    NameFinder = LocalNameFinder

    FunctionGen = None

    ClassGen = None

    def visitModule(self, node):

        self.scopes = self.parseSymbols(node)

        self.scope = self.scopes[node]

        self.emit('SET_LINENO', 0)

        if node.doc:

            self.emit('LOAD_CONST', node.doc)

            self.storeName('__doc__')

        lnf = walk(node.node, self.NameFinder(), verbose=0)

        self.locals.push(lnf.getLocals())

        self.visit(node.node)

        self.emit('LOAD_CONST', None)

        self.emit('RETURN_VALUE')

    def visitExpression(self, node):

        self.set_lineno(node)

        self.scopes = self.parseSymbols(node)

        self.scope = self.scopes[node]

        self.visit(node.node)

        self.emit('RETURN_VALUE')

    def visitFunction(self, node):

        self._visitFuncOrLambda(node, isLambda=0)

        if node.doc:

            self.setDocstring(node.doc)

        self.storeName(node.name)

    def visitLambda(self, node):

        self._visitFuncOrLambda(node, isLambda=1)

    def _visitFuncOrLambda(self, node, isLambda=0):

        if not isLambda and node.decorators:

            for decorator in node.decorators.nodes:

                self.visit(decorator)

            ndecorators = len(node.decorators.nodes)

        else:

            ndecorators = 0

        gen = self.FunctionGen(node, self.scopes, isLambda,

                               self.class_name, self.get_module())

        walk(node.code, gen)

        gen.finish()

        self.set_lineno(node)

        for default in node.defaults:

            self.visit(default)

        self._makeClosure(gen, len(node.defaults))

        for i in range(ndecorators):

            self.emit('CALL_FUNCTION', 1)

    def visitClass(self, node):

        gen = self.ClassGen(node, self.scopes,

                            self.get_module())

        walk(node.code, gen)

        gen.finish()

        self.set_lineno(node)

        self.emit('LOAD_CONST', node.name)

        for base in node.bases:

            self.visit(base)

        self.emit('BUILD_TUPLE', len(node.bases))

        self._makeClosure(gen, 0)

        self.emit('CALL_FUNCTION', 0)

        self.emit('BUILD_CLASS')

        self.storeName(node.name)

    # The rest are standard visitor methods

    # The next few implement control-flow statements

    def visitIf(self, node):

        end = self.newBlock()

        numtests = len(node.tests)

        for i in range(numtests):

            test, suite = node.tests[i]

            if is_constant_false(test):

                # XXX will need to check generator stuff here

                continue

            self.set_lineno(test)

            self.visit(test)

            nextTest = self.newBlock()

            self.emit('JUMP_IF_FALSE', nextTest)

            self.nextBlock()

            self.emit('POP_TOP')

            self.visit(suite)

            self.emit('JUMP_FORWARD', end)

            self.startBlock(nextTest)

            self.emit('POP_TOP')

        if node.else_:

            self.visit(node.else_)

        self.nextBlock(end)

    def visitWhile(self, node):

        self.set_lineno(node)

        loop = self.newBlock()

        else_ = self.newBlock()

        after = self.newBlock()

        self.emit('SETUP_LOOP', after)

        self.nextBlock(loop)

        self.setups.push((LOOP, loop))

        self.set_lineno(node, force=True)

        self.visit(node.test)

        self.emit('JUMP_IF_FALSE', else_ or after)

        self.nextBlock()

        self.emit('POP_TOP')

        self.visit(node.body)

        self.emit('JUMP_ABSOLUTE', loop)

        self.startBlock(else_) # or just the POPs if not else clause

        self.emit('POP_TOP')

        self.emit('POP_BLOCK')

        self.setups.pop()

        if node.else_:

            self.visit(node.else_)

        self.nextBlock(after)

    def visitFor(self, node):

        start = self.newBlock()

        anchor = self.newBlock()

        after = self.newBlock()

        self.setups.push((LOOP, start))

        self.set_lineno(node)

        self.emit('SETUP_LOOP', after)

        self.visit(node.list)

        self.emit('GET_ITER')

        self.nextBlock(start)

        self.set_lineno(node, force=1)

        self.emit('FOR_ITER', anchor)

        self.visit(node.assign)

        self.visit(node.body)

        self.emit('JUMP_ABSOLUTE', start)

        self.nextBlock(anchor)

        self.emit('POP_BLOCK')

        self.setups.pop()

        if node.else_:

            self.visit(node.else_)

        self.nextBlock(after)

    def visitBreak(self, node):

        if not self.setups:

            raise SyntaxError, "'break' outside loop (%s, %d)" % \

                  (node.filename, node.lineno)

        self.set_lineno(node)

        self.emit('BREAK_LOOP')

    def visitContinue(self, node):

        if not self.setups:

            raise SyntaxError, "'continue' outside loop (%s, %d)" % \

                  (node.filename, node.lineno)

        kind, block = self.setups.top()

        if kind == LOOP:

            self.set_lineno(node)

            self.emit('JUMP_ABSOLUTE', block)

            self.nextBlock()

        elif kind == EXCEPT or kind == TRY_FINALLY:

            self.set_lineno(node)

            # find the block that starts the loop

            top = len(self.setups)

            while top > 0:

                top = top - 1

                kind, loop_block = self.setups[top]

                if kind == LOOP:

                    break

            if kind != LOOP:

                raise SyntaxError, "'continue' outside loop (%s, %d)" % \

                      (node.filename, node.lineno)

            self.emit('CONTINUE_LOOP', loop_block)

            self.nextBlock()

        elif kind == END_FINALLY:

            msg = "'continue' not allowed inside 'finally' clause (%s, %d)"

            raise SyntaxError, msg % (node.filename, node.lineno)

    def visitTest(self, node, jump):

        end = self.newBlock()

        for child in node.nodes[:-1]:

            self.visit(child)

            self.emit(jump, end)

            self.nextBlock()

            self.emit('POP_TOP')

        self.visit(node.nodes[-1])

        self.nextBlock(end)

    def visitAnd(self, node):