#!/www/python/bin/python

"""
Compile a PTL template.

First the tokens "template" are replaced with "def".  Next, the file is
parsered into an parse tree.  This tree is converted into a modified AST.
It is during this state that the semantics are modified by adding extra
nodes to the tree.  Finally bytecode is generated using the compiler
package.

Note that script/module requires the compiler package.
"""

# created 2000/07/28, nas

__revision__ = "$Id: ptl_compile.py,v 1.17 2001/08/29 17:52:35 gward Exp $"

import sys
import os
import string
import symbol
import token
import parser
import re

from compiler import pycodegen, visitor, transformer, walk
from compiler import ast
from compiler.consts import OP_ASSIGN

# magic names inserted into the code

# try to figure out what package we are in
_package = string.split(__name__, ".")[:-1]
IO_MODULE = string.join(_package + ["TemplateIO"], ".")
IO_CLASS = "TemplateIO"
IO_INSTANCE = "__output"
TEMPLATE_PREFIX = "__template__"

class TemplateTransformer(transformer.Transformer):

    def __init__ (self, *args, **kwargs):
        transformer.Transformer.__init__(self, *args, **kwargs)
        self.__in_template = [] # stack, 1 if in template else 0

    def file_input(self, nodelist):
        # Add a "from IO_MODULE import IO_CLASS" statement to the
        # beginning of the module.
        doc = None # self.get_docstring(nodelist, symbol.file_input)
        if sys.hexversion >= 0x20000b1:
            # imports work differently since Python 2.0b1
            imp = ast.From(IO_MODULE, [(IO_CLASS, None)])
        else:
            imp = ast.From(IO_MODULE, [IO_CLASS])

        # Add an IO_INSTANCE binding for module level expressions (like
        # doc strings).  This instance will not be returned.
        klass = ast.Name(IO_CLASS)
        instance = ast.CallFunc(klass, [])
        assign_name = ast.AssName(IO_INSTANCE, OP_ASSIGN)
        assign = ast.Assign([assign_name], instance)

        stmts = [ imp, assign ]

        for node in nodelist:
            if node[0] != token.ENDMARKER and node[0] != token.NEWLINE:
                self.com_append_stmt(stmts, node)

        return ast.Module(doc, ast.Stmt(stmts))

    def funcdef(self, nodelist):
        lineno = nodelist[1][2]
        name     = nodelist[1][1]
        args     = nodelist[2][2]

        if not name.startswith(TEMPLATE_PREFIX):
            # just a normal function, let base class handle it 
            self.__in_template.append(0)
            n = transformer.Transformer.funcdef(self, nodelist)

        else:
            self.__in_template.append(1)

            name = name[len(TEMPLATE_PREFIX):]

            # Add "IO_INSTANCE = IO_CLASS()" statement at the beginning of
            # the function and a "return IO_INSTANCE" at the end.
            if args[0] == symbol.varargslist:
                names, defaults, flags = self.com_arglist(args[1:])
            else:
                names = defaults = ()
                flags = 0
            doc = None # self.get_docstring(nodelist[4])

            # code for function
            code = self.com_node(nodelist[4])

            # create an instance, assign to IO_INSTANCE
            klass = ast.Name(IO_CLASS)
            instance = ast.CallFunc(klass, [])
            assign_name = ast.AssName(IO_INSTANCE, OP_ASSIGN)
            assign = ast.Assign([assign_name], instance)

            # return the str() of the instance
            func = ast.Getattr(ast.Name(IO_INSTANCE), "__str__")
            ret = ast.Return(ast.CallFunc(func, []))

            # wrap original function code
            code = ast.Stmt([assign, code, ret])

            n = ast.Function(name, names, defaults, flags, doc, code)
            n.lineno = lineno

        self.__in_template.pop()
        return n

    def expr_stmt(self, nodelist):
        if not self.__in_template or not self.__in_template[-1]:
            return transformer.Transformer.expr_stmt(self, nodelist)

        # Instead of discarding objects on the stack, call
        # "IO_INSTANCE.write(obj)".
        exprNode = self.com_node(nodelist[-1])
        if len(nodelist) == 1:
            lval = ast.Name(IO_INSTANCE)
            n = ast.AugAssign(lval, '+=', exprNode)
            if hasattr(exprNode, 'lineno'):
                n.lineno = exprNode.lineno
        elif nodelist[1][0] == token.EQUAL:
            nodes = [ ]
            for i in range(0, len(nodelist) - 2, 2):
                nodes.append(self.com_assign(nodelist[i], OP_ASSIGN))
            n = ast.Assign(nodes, exprNode)
            n.lineno = nodelist[1][2]
        else:
            lval = self.com_augassign(nodelist[0])
            op = self.com_augassign_op(nodelist[1])
            n = ast.AugAssign(lval, op[1], exprNode)
            n.lineno = op[2]
        return n


_template_re = re.compile(r"^([ \t]*)template([ \t]+)", re.MULTILINE)

def translate_tokens(buf, filename):
    """
    Since we can't modify the parser in the builtin parser module we
    must do token translation here.  Luckily it does not affect line
    numbers.
   
    template foo(...): -> def __template__foo(...):
   
    XXX This parser is too stupid.  For example, it doesn't understand
    triple quoted strings.
    """
    global _def_re, _template_re

    # change template to def
    buf = _template_re.sub(r"\1def\2%s" % TEMPLATE_PREFIX, buf)

    # compile() and parsermodule don't accept certain modules if they are
    # missing a trailing newline.  The Python interpreter seems to add a
    # newline when importing modules.  We will match that its behavior.
    if buf and buf[-1] != '\n':
        buf += "\n"

    return buf


def parse(buf, filename='<string>'):
    buf = translate_tokens(buf, filename)
    try:
        return TemplateTransformer().parsesuite(buf)
    except (parser.ParserError, SyntaxError):
        # Let Python compiler generate a better error message
        import __builtin__, sys
        try:
            __builtin__.compile(buf, filename, 'exec')
        except SyntaxError, value:
            # Python does not use filename in the exception but stupidly
            # uses "<string>" instead.  Patch it up ourself.
            lineno = value.lineno
            offset = value.offset
            text = value.text
            raise SyntaxError(str(value), (filename, lineno, offset, text))
        raise # hmm, the Python compiler didn't raise an exception!?


PTL_EXT = ".ptl"
PTLC_EXT = ".ptlc"
if sys.hexversion >= 0x20100b1:
    PTLC_MAGIC = "PTLC\000\005"
elif sys.hexversion >= 0x20000b1:
    PTLC_MAGIC = "PTLC\000\004"
else:
    PTLC_MAGIC = "PTLC\000\003"


class Template(pycodegen.Module):

    def compile(self):
        ast = parse(self.source, self.filename)
        root, filename = os.path.split(self.filename)
        gen = pycodegen.ModuleCodeGenerator(filename)
        walk(ast, gen, 1)
        self.code = gen.getCode()

    def dump(self, f):
        import marshal
        import stat
        import struct

        f.write(PTLC_MAGIC)
        mtime = os.stat(self.filename)[stat.ST_MTIME]
        marshal.dump(mtime, f)
        marshal.dump(self.code, f)


def compile_template(input, filename, output=None):
    """compile_template(input, filename, output=None) -> code
    
    Compile an open file.  If output is not None then the code is written
    with the magic template header.  The code object is returned.
    """
    buf = input.read()
    template = Template(buf, filename)
    template.compile()
    if output:
        template.dump(output)
    return template.code

def compile(inputname, outputname):
    """compile(inputname, outputname)
    
    Compile a template file.  The new template is writen to outputname.
    """
    input = open(inputname)
    output = open(outputname, "wb")
    try:
        compile_template(input, inputname, output)
    except:
        # don't leave a corrupt .ptlc file around
        output.close()
        os.unlink(outputname)
        raise

def main():
    args = sys.argv[1:]
    if not args:
        print "no files to compile"
    else:
        for filename in args:
            path, ext = os.path.splitext(filename)
            compile(filename, path + PTLC_EXT)

if __name__ == "__main__":
    main()
