LINUX.ORG.RU

Замена узлов AST дерева и конвертация дерева в код?

 


0

2

Сабж. Есть некоторый фрагмент питоньего кода (строка), есть словарь замен. Надо в строке кода некоторые переменные поменять согласно словарю на то что в словаре и получить код после преобразования.

Словарь имеет вид

table = {'A':1.2, 'B':3.14, 'X':'Y'}
подразумевается что переменная X должна быть заменена на переменную Y и т.д.

Вроде я могу при помощи модуля ast конвертнуть код в дерево и провести в нем замены. Но как мне обратно собрать код? Все что я попытался сделать с codegen, sourcecodegen пока не работает.

Код для теста

#!/usr/bin/python
import ast, sourcecodegen

table = {'A':123, 'X':'Y'}

class Repl(ast.NodeTransformer):
    def visit_Name(self, node):
        if node.id in table: return ast.parse(str(table[node.id]))

E = 'sin(X)*A'

P = ast.parse(E, 'eval')
R = Repl().visit(P)

#print astor.codegen.to_source(R)
print sourcecodegen.generate_code(R)

ЗЫ Хотелось бы что бы решение было на основе стандартных модулей, не требующих установки чего то дополнительного

★★

Disclaimer: питон я не знаю.

Вроде я могу при помощи модуля ast конвертнуть код в дерево и провести в нем замены

Ты это делаешь неправильно. Во-первых, visit_Name должен всегда возвращать значение, а не только если происходит замена. Во-вторых, как минимум в зависимости от типа замены нужно заменять текущий узел на узлы разных типов. Т.е. если замена — строка, ты заменяешь на ast.Name, если число — то на ast.Numb и т.д.

ЗЫ Хотелось бы что бы решение было на основе стандартных модулей, не требующих установки чего то дополнительного

Судя по тому, что я нагуглил, в стандартной библиотеке для этого ничего нет, поэтому использовал astor:

#!/usr/bin/python

import ast
import astor

class Repl(ast.NodeTransformer):
    def __init__(self, table):
        super(Repl, self).__init__()

        self.__table = table

    def visit_Name(self, node):
        if node.id in self.__table:
            replacement = self.__table[node.id]
            replacement_type = type(replacement)

            new_node = node

            if replacement_type == str:
                new_node = ast.Name(replacement, node.ctx)

            if replacement_type == int or replacement_type == float:
                new_node = ast.Num(replacement)

            # TODO: handle more types of replament values

            return ast.copy_location(new_node, node)
        else:
            return node

expr = 'sin(X) * A'
table = { 'X': 'Y', 'A': 123 }

parsed_ast = ast.parse(expr)

transformed_ast = Repl(table).visit(parsed_ast)
print astor.to_source(transformed_ast)
theNamelessOne ★★★★★ ()