Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
Код для модификации AST с доклада на PyCon Russia 2015
# -*- 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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment