/qpy/compile_jython.py

"""
open/DurusWorks/qpy/compile_jython.py

"""
from java.io import ByteArrayInputStream
from java.io import ByteArrayOutputStream
from org.python.antlr import AST
from org.python.compiler import Module
from org.python.core import imp
from org.python.core import ParserFacade, CompileMode, CompilerFlags
from os.path import splitext
from pprint import pprint
from qpy.translate import translate_tokens
from _ast import astlist
import py_compile
import re
import symbol
import sys
import sys
import token
import __builtin__ as builtins


def get_parse_tree(source, source_name=None):
    if source_name is None:
        mode = 'eval'
    else:
        mode = 'exec'
    node = ParserFacade.parse(
          ByteArrayInputStream(source),
          getattr(CompileMode, mode),
          source_name or '',
          CompilerFlags())
    return node

def tree(node):
    if isinstance(node, AST):
        result = [node.__class__.__name__]
        result.append(
            [(name, tree(getattr(node, name))) for name in node._fields])
        result.append(
            [(name, tree(getattr(node, name))) for name in node._attributes])
        if hasattr(node, '__iter__'):
            for child in node:
                result.append((type(child), tree(child)))
        return result
    if node is None:
        return None
    if type(node) in (str, int, float):
        return node
    if hasattr(node, '__iter__'):
        return [node.__class__.__name__] + [tree(x) for x in node]
    else:
        return (type(node), node)

def ptree(node):
    pprint(tree(node))

def children(node):
    if isinstance(node, AST):
        result = []
        for name in node._fields:
            if name == 'ctx':
                continue
            value = getattr(node, name)
            if get_node_type(node) == 'List' or isinstance(value, astlist):
                #print name, value
                #ptree(value)
                for x in value:
                    result.append(x)
            else:
                result.append(value)
        return result
    if hasattr(node, '__iter__'):
        return list(node)
    else:
        return []

def get_node_type(node):
    return node.__class__.__name__


class Transform (object):

    def __init__(self, tree):

        self.tree = tree
        self.template_type_stack = [None]
        self.traverse_node(self.tree)
        self.line_number = 1
        self.rationalize_line_numbers(self.tree)

    def rationalize_line_numbers(self, node):
        if hasattr(node, 'lineno'):
            if not node.lineno or node.lineno < self.line_number:
                node.lineno = self.line_number
            else:
                self.line_number = node.lineno
        for n in children(node):
            self.rationalize_line_numbers(n)

    def traverse_node(self, node):
        node_type = get_node_type(node)
        #print node_type, node, children(node)
        if node_type == 'FunctionDef':
            function_name = node.name
            if function_name.endswith('__xml_template__'):
                self.template_type_stack.append('xml')
            elif function_name.endswith('__str_template__'):
                self.template_type_stack.append('str')
            else:
                self.template_type_stack.append(None)
        # Quote literals in xml templates first.
        for child in children(node):
            if self.template_type_stack[-1] == 'xml' and 'Str' == get_node_type(child):
                if node_type in ('Expr', 'Attribute') and child is node.value:
                    xml_quote_call = get_parse_tree('_qpy_xml("X")').body
                    xml_quote_call.args[0] = child
                    node.value = xml_quote_call
                elif node_type == 'BinOp' and child is node.left:
                    xml_quote_call = get_parse_tree('_qpy_xml("X")').body
                    xml_quote_call.args[0] = child
                    node.left = xml_quote_call
                elif node_type == 'BinOp' and child is node.right:
                    xml_quote_call = get_parse_tree('_qpy_xml("X")').body
                    xml_quote_call.args[0] = child
                    node.right = xml_quote_call
                elif node_type == 'Call':
                    for j, arg in enumerate(node.args):
                        if get_node_type(arg) == 'Str':
                            xml_quote_call = get_parse_tree('_qpy_xml("X")').body
                            xml_quote_call.args[0] = arg
                            node.args[j] = xml_quote_call
                elif node_type == 'BinOp':
                    for j, arg in enumerate(node.args):
                        if get_node_type(arg) == 'Str':
                            xml_quote_call = get_parse_tree('_qpy_xml("X")').body
                            xml_quote_call.args[0] = arg
                            node.args[j] = xml_quote_call
                elif node_type == 'Dict':
                    for j, arg in enumerate(node.keys):
                        if get_node_type(arg) == 'Str':
                            xml_quote_call = get_parse_tree('_qpy_xml("X")').body
                            xml_quote_call.args[0] = arg
                            node.keys[j] = xml_quote_call
                    for j, arg in enumerate(node.values):
                        if get_node_type(arg) == 'Str':
                            xml_quote_call = get_parse_tree('_qpy_xml("X")').body
                            xml_quote_call.args[0] = arg
                            node.values[j] = xml_quote_call
                elif node_type in ('List', 'Tuple'):
                    for j, arg in enumerate(node.elts):
                        if get_node_type(arg) == 'Str':
                            xml_quote_call = get_parse_tree('_qpy_xml("X")').body
                            xml_quote_call.args[0] = arg
                            node.elts[j] = xml_quote_call
                else:
                    print type(node)
                    ptree(node)
                    assert 0
            self.traverse_node(child)
        if node_type == 'Module':
            for index, child in enumerate(children(node)):
                child_type = get_node_type(child)
                if index == 0 and child_type == 'Expr' and get_node_type(child.value) == 'Str':
                    continue
                if child_type == 'ImportFrom' and child.module == '__future__':
                    continue
                module_start = get_parse_tree(
                    'from qpy.quoted import join_xml as _qpy_join_xml, '
                    'join_str as _qpy_join_str, xml as _qpy_xml', 'fake').body[0]
                node.body.insert(index, module_start)
                break
        elif self.template_type_stack[-1] is None:
            pass # We're not in a template, so we're done.
        elif node_type == 'Expr':
            # Wrap this expression statement in a qpy_append call.
            new_expr = get_parse_tree('qpy_append(X)', 'fake').body[0]
            new_expr.value.args[0] = node.value
            node.value = new_expr.value
        elif node_type == "FunctionDef":
            node.name = node.name[:-16]
            template_start = get_parse_tree('qpy_accumulation=[];qpy_append=qpy_accumulation.append', 'fake').body
            node.body.insert(0, template_start[0])
            node.body.insert(1, template_start[1])
            if self.template_type_stack[-1] == 'xml':
                template_end = get_parse_tree('return _qpy_join_xml(qpy_accumulation)', 'fake').body[0]
            else:
                template_end = get_parse_tree('return _qpy_join_str(qpy_accumulation)', 'fake').body[0]
            node.body.append(template_end)
        if node_type == 'FunctionDef':
            self.template_type_stack.pop()

    def get_result(self):
        return self.tree


PYC = "$py.class"

def get_code(source, source_name):
    translated_source = translate_tokens(source)
    node = get_parse_tree(translated_source, source_name)
    transformed = Transform(node).get_result()
    code = ByteArrayOutputStream()
    transformed = node
    Module.compile(transformed, code, source_name +"$py", source_name, True, False, None, imp.NO_MTIME)
    return code

def compile(source_name):
    source = open(source_name).read()
    code = get_code(source, source_name)
    class_file_name = splitext(source_name)[0] + PYC
    class_file = open(class_file_name, 'wb')
    class_file.write(code.toByteArray())
    class_file.close()



if __name__ == '__main__':
    import qpy.__main__