Измените несколько узлов с помощью python ast.NodeTransformer - PullRequest
Новая
       27

Измените несколько узлов с помощью python ast.NodeTransformer

0 голосов
/ 09 июля 2020

У меня есть входной исходный код, подобный этому 

def foo(my_input):
  return my_input + 42

, и я хочу, чтобы он преобразовался вот так 

def method_name(arg0):
  return my_input + 42

Преобразователь узла ast для этой цели записан так. 

class MyRenamer(ast.NodeTransformer):

  def __init__(self):
    self._arg_count = 0

  def visit_FunctionDef(self, node):
    node.name = "method_name"
    return node

  def visit_arg(self, node):
    node.arg = "arg_{}".format(self._arg_count)
    self._arg_count += 1
    return node

но когда я вызываю вышеуказанный трансформатор вот так. 

node = ast.parse(code)
renamer = MyRenamer()
node2 = renamer.visit(node)
print(astor.to_source(node2))

Я получаю 

def method_name(my_input):
    return my_input + 42

Здесь аргумент функции не был изменен .

1 Ответ

1 голос
/ 13 июля 2020

Посетителю необходимо пройти AST, посетив всех дочерних узлов текущего посещаемого узла. Метод generic_visit () делает это за вас, но вы должны вызывать его при каждом методе visit _... или, по крайней мере, в тех случаях, когда дети являются возможными. 

import ast
import astor

class MyRenamer(ast.NodeTransformer):

  def __init__(self):
    self._arg_count = 0

  def visit_FunctionDef(self, node):
    node.name = "method_name"
    self.generic_visit(node)
    return node

  def visit_arg(self, node):
    node.arg = "arg_{}".format(self._arg_count)
    self._arg_count += 1
    self.generic_visit(node)
    return node

code = """
def foo(my_input):
  return my_input + 42
"""

node = ast.parse(code)
renamer = MyRenamer()
node2 = renamer.visit(node)
print(astor.to_source(node2))


def method_name(arg_0):
    return my_input + 42

Это дает ожидаемый результат, но в более широком контексте все функции будут переименованы в "method_name", что может быть нежелательно. И еще есть идентификатор в теле функции, который, по-видимому, тоже нужно переименовать.

...