Код для модификации AST с доклада на PyCon Russia 2015

aster.py

# -*- coding: utf-8 -*-

import ast
from ast import (
    Assign,
    Subscript,
    Name,
    Index,
    NodeTransformer,
    Tuple,
    copy_location,
)
import astunparse
import astor
from astmonkey import visitors


def make_tree(fn):
    """
    Преобразуем исходный код в AST.
    """
    with open(fn, 'r') as fp:
        return ast.parse(fp.read())


def check_structure(node):
    """
    Фильтруем узлы дерева, проверяя его структуру и типы содержимого
    атрибутов.
    """
    if not isinstance(node, Assign):
        return False

    if not isinstance(node.value, Subscript):
        return False

    if not isinstance(node.value.value, Name):
        return False

    if not isinstance(node.value.slice, Index):
        return False

    return len(node.targets) == 1


class Compressor(NodeTransformer):
    def same_source(self, a, b):
        """
        Смотрим на контекст - нас интересуют только узлы присвоения,
        у которых источником данных является один и тот же массив.
        """
        return a and a.value.value.id == b.value.value.id

    def compress(self, assigns, contents):
        """
        Вносим изменения: последовательные присвоения переменным
        значений из одного и того же массива преобразуем в более
        компактное представление.
        """
        targets, values = [], []

        if not assigns:
            return

        if len(assigns) == 1:
            compressed = assigns[0]
        else:
            for node in assigns:
                targets.append(node.targets[0])
                values.append(node.value)

            compressed = Assign(
                targets=[Tuple(elts=targets)],
                value=Tuple(elts=values),
            )
            copy_location(compressed, assigns[0])
        contents.append(compressed)


    def visit(self, node):
        """
        "Все вместе!"

        Соединяем вместе раздельные элементы:

        * Ищем узлы, которые ищут другие узлы.
        * Проверяем, что найденные узлы и их соседи отвечают
        нашим условиям.
        * Схлопываем похожие узлы с выражениями присвоения в один узел.
        * Выгружаем все обратно в исходный код.
        """
        # Посещаем все дочерние узлы.
        self.generic_visit(node)

        # Нас не интересует узлы, которые не содержат выражений.
        if not hasattr(node, 'body'):
            return node

        last_found = None
        new_contents = []
        candidates = []

        # Последовательно перебираем все дочерние узлы.
        for child in node.body:
            # Проверяем, что данный узел является присвоением
            # одной переменной значения из явно заданного элемента
            # массива.
            if check_structure(child):
                if not self.same_source(last_found, child):
                    self.compress(candidates, new_contents)
                    candidates = []

                candidates.append(child)
                last_found = child
            else:
                self.compress(candidates, new_contents)
                new_contents.append(child)
                candidates, last_found = [], None

        self.compress(candidates, new_contents)
        node.body = new_contents
        return node


tree = make_tree('sample.py')
modified_tree = Compressor().visit(tree)

print ''
print '#### astunparse ####'
print astunparse.unparse(modified_tree)

print ''
print '#### astor ####'
print astor.to_source(modified_tree)

print ''
print '#### astmonkey ####'
print visitors.to_source(modified_tree)