<commit_before>from ast import NodeTransformer, iter_fields, copy_location
from itertools import chain

import astor

from .nodes import String, Tuple, Symbol, List, Number, Placeholder
from .nodes import NodeVisitor
from .compat import ast as py, texttype
from .checker import split_args
from .constant import HTML_ELEMENTS, SELF_CLOSING_ELEMENTS


def _write(value):
    return py.Expr(py.Call(
        py.Attribute(py.Name('buf', py.Load()), 'write', py.Load()),
        [value], [], None, None,
    ))


def _write_str(value):
    return _write(py.Str(value))


def _ctx_load(name):
    return py.Attribute(py.Name('ctx', py.Load()), name, py.Load())


def _ctx_store(name):
    return py.Attribute(py.Name('ctx', py.Load()), name, py.Store())


def _buf_push():
    return py.Expr(py.Call(
        py.Attribute(py.Name('buf', py.Load()), 'push', py.Load()),
        [], [], None, None,
    ))


def _buf_pop():
    return py.Call(py.Attribute(py.Name('buf', py.Load()),
                                'pop', py.Load()),
                   [], [], None, None)


def _returns_output_type(node):
    # FIXME: temporary hack before we will have working types system
    return (
        isinstance(node, Tuple) and
        (node.values[0].name in {'each', 'if', 'join'} or
         node.values[0].name in HTML_ELEMENTS or
         node.values[0].ns)
    )


def _yield_writes(node):
    if _returns_output_type(node):
        for item in compile_(node, True):
            yield item
    else:
        for item in compile_(node, True):
            yield _write(item)


class _PlaceholdersExtractor(NodeVisitor):

    def __init__(self):
        # using list to preserve placeholders order for the tests
        self.placeholders = []

    def visit_placeholder(self, node):
        if node.name not in self.placeholders:
            self.placeholders.append(node.name)


_cls_eq = lambda i, name: i.__class__.__name__ == name


def _node_copy(func):
    def wrapper(self, node):
        node = self.generic_visit(node)
        node_cls = type(node)
        new_node = node_cls(*[value for _, value in iter_fields(node)])
        copy_location(new_node, node)
        func(self, new_node)
        return new_node
    return wrapper


class _Optimizer(NodeTransformer):

    def _paste(self, body):
        chunks = []
        for item in body:
            if (
                _cls_eq(item, 'Expr') and
                _cls_eq(item.value, 'Call') and
                _cls_eq(item.value.func, 'Attribute') and
                _cls_eq(item.value.func.value, 'Name') and
                item.value.func.value.id == 'buf' and
                item.value.func.attr == 'write' and
                len(item.value.args) == 1 and
                (_cls_eq(item.value.args[0], 'Str') or
                 _cls_eq(item.value.args[0], 'Num'))
            ):
                if _cls_eq(item.value.args[0], 'Str'):
                    chunks.append(item.value.args[0].s)
                else:
                    chunks.append(texttype(item.value.args[0].n))
            else:
                if chunks:
                    yield _write_str(u''.join(chunks))
                    del chunks[:]
                yield item
        if chunks:
            yield _write_str(u''.join(chunks))

    @_node_copy
    def visit_Module(self, node):
        node.body = list(self._paste(node.body))

    @_node_copy
    def visit_FunctionDef(self, node):
        node.body = list(self._paste(node.body))

    @_node_copy
    def visit_If(self, node):
        node.body = list(self._paste(node.body))
        node.orelse = list(self._paste(node.orelse))

    @_node_copy
    def visit_For(self, node):
        node.body = list(self._paste(node.body))


def _copy_location(func):
    def wrapper(node, as_statement):
        for item in func(node, as_statement):
            if not getattr(item, 'lineno', None):
                item.lineno = node.location.start.line
            yield item
    return wrapper


@_copy_location
def compile_(node, as_statement):
    if isinstance(node, Tuple):
        sym, args = node.values[0], node.values[1:]
        pos_args, kw_args = split_args(args)
        if sym.name == 'def':
            assert as_statement
            name_sym, body = pos_args
            visitor = _PlaceholdersExtractor()
            visitor.visit(body)
            py_args = [py.arg(p) for p in visitor.placeholders]
            yield py.FunctionDef(name_sym.name,
                                 py.arguments(py_args, None, None, []),
                                 list(compile_(body, as_statement)), [])
<commit_msg>Removed location info from compiled code, to replace it later with more complete solution.<commit_after>from ast import NodeTransformer, iter_fields, copy_location
from itertools import chain

import astor

from .nodes import String, Tuple, Symbol, List, Number, Placeholder
from .nodes import NodeVisitor
from .compat import ast as py, texttype
from .checker import split_args
from .constant import HTML_ELEMENTS, SELF_CLOSING_ELEMENTS


def _write(value):
    return py.Expr(py.Call(
        py.Attribute(py.Name('buf', py.Load()), 'write', py.Load()),
        [value], [], None, None,
    ))


def _write_str(value):
    return _write(py.Str(value))


def _ctx_load(name):
    return py.Attribute(py.Name('ctx', py.Load()), name, py.Load())


def _ctx_store(name):
    return py.Attribute(py.Name('ctx', py.Load()), name, py.Store())


def _buf_push():
    return py.Expr(py.Call(
        py.Attribute(py.Name('buf', py.Load()), 'push', py.Load()),
        [], [], None, None,
    ))


def _buf_pop():
    return py.Call(py.Attribute(py.Name('buf', py.Load()),
                                'pop', py.Load()),
                   [], [], None, None)


def _returns_output_type(node):
    # FIXME: temporary hack before we will have working types system
    return (
        isinstance(node, Tuple) and
        (node.values[0].name in {'each', 'if', 'join'} or
         node.values[0].name in HTML_ELEMENTS or
         node.values[0].ns)
    )


def _yield_writes(node):
    if _returns_output_type(node):
        for item in compile_(node, True):
            yield item
    else:
        for item in compile_(node, True):
            yield _write(item)


class _PlaceholdersExtractor(NodeVisitor):

    def __init__(self):
        # using list to preserve placeholders order for the tests
        self.placeholders = []

    def visit_placeholder(self, node):
        if node.name not in self.placeholders:
            self.placeholders.append(node.name)


_cls_eq = lambda i, name: i.__class__.__name__ == name


def _node_copy(func):
    def wrapper(self, node):
        node = self.generic_visit(node)
        node_cls = type(node)
        new_node = node_cls(*[value for _, value in iter_fields(node)])
        copy_location(new_node, node)
        func(self, new_node)
        return new_node
    return wrapper


class _Optimizer(NodeTransformer):

    def _paste(self, body):
        chunks = []
        for item in body:
            if (
                _cls_eq(item, 'Expr') and
                _cls_eq(item.value, 'Call') and
                _cls_eq(item.value.func, 'Attribute') and
                _cls_eq(item.value.func.value, 'Name') and
                item.value.func.value.id == 'buf' and
                item.value.func.attr == 'write' and
                len(item.value.args) == 1 and
                (_cls_eq(item.value.args[0], 'Str') or
                 _cls_eq(item.value.args[0], 'Num'))
            ):
                if _cls_eq(item.value.args[0], 'Str'):
                    chunks.append(item.value.args[0].s)
                else:
                    chunks.append(texttype(item.value.args[0].n))
            else:
                if chunks:
                    yield _write_str(u''.join(chunks))
                    del chunks[:]
                yield item
        if chunks:
            yield _write_str(u''.join(chunks))

    @_node_copy
    def visit_Module(self, node):
        node.body = list(self._paste(node.body))

    @_node_copy
    def visit_FunctionDef(self, node):
        node.body = list(self._paste(node.body))

    @_node_copy
    def visit_If(self, node):
        node.body = list(self._paste(node.body))
        node.orelse = list(self._paste(node.orelse))

    @_node_copy
    def visit_For(self, node):
        node.body = list(self._paste(node.body))


def compile_(node, as_statement):
    if isinstance(node, Tuple):
        sym, args = node.values[0], node.values[1:]
        pos_args, kw_args = split_args(args)
        if sym.name == 'def':
            assert as_statement
            name_sym, body = pos_args
            visitor = _PlaceholdersExtractor()
            visitor.visit(body)
            py_args = [py.arg(p) for p in visitor.placeholders]
            yield py.FunctionDef(name_sym.name,
                                 py.arguments(py_args, None, None, []),
                                 list(compile_(body, as_statement)), [])

{'fraction_non_alphanumeric': 0.1219917932793612, 'fraction_numerical': 0.0024398358655872243, 'mean_word_length': 2.1247401247401245, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7415270', 'n_tokens_mistral': 2815, 'n_tokens_neox': 2730, 'n_words': 647}