Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
Код для модификации исходного кода с помощью redbaron, PyCon Russia 2015
# -*- coding: utf-8 -*-
from redbaron import RedBaron, AssignmentNode, TupleNode
def same_source(a, b):
"""
Фильтр узлов дерева.
"""
return a.value[0].value == b.value[0].value
def compress(assigns):
"""
Подготовка изменений с учетом констекста.
"""
v = u'%s = %s' % (
', '.join([n.target.dumps() for n in assigns]),
', '.join([n.value.dumps() for n in assigns])
)
return v
def strip_duplicate_endls(root):
"""
Убираем лишние переводы строк.
"""
nl = root.filter(lambda n: n.type == 'endl' and n.next)
nl = [
n
for n in nl
if n.next.type == 'endl' and n in root
]
for endl in nl:
root.remove(endl)
def group_assigns(start):
"""
Схлопываем похожие между собой присвоения.
"""
if start not in start.parent:
return
found, node, parent = [], start.next, start.parent
while node:
if node.type not in ['endl', 'assignment']:
break
if node.type == 'assignment' and same_source(node, start):
found.append(node)
node = node.next
if found:
start.replace(compress([start] + found))
# Удаляем узлы, которые мы "схлопнули"
for node in found:
if node in parent:
parent.remove(node)
strip_duplicate_endls(parent)
# --------------------------
# "И ещё раз, все вместе..."
# --------------------------
with open('sample.py', 'r') as fp:
red = RedBaron(fp.read())
assigns = red.find_all(
'assign',
target=lambda t: len(t) == 1,
value=lambda v: len(v) == 2 and v[1].type == 'getitem',
)
for node in assigns:
group_assigns(node)
print red.dumps()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment