使用 Python 重命名 Python 代码中的变量

Question

>我有一些Python代码作为字符串，

string = """
import numpy as nmp
y = 5
def f(x):
return nmp.sum(x) + y
x = 1
print(f(x))
"""

使用Python，我想重命名 导入到nmp到np， 函数参数xX(但不是x = 1)，以及yY.

为此，我必须能够识别各个变量的每种用法。我想我必须使用 ast 或 libcst 之一，但我不太确定。

有什么提示吗？

Answer 1

下面的代码将源代码作为ast.AST对象遍历，根据提供的映射更新名称。此映射列出了要更改的对象的新名称，并指定了应禁止替换的任何范围：

import ast
def walk(tree, mp, scope=['main']):
for i in tree._fields:
if not isinstance(a:=getattr(tree, i), list): 
if a in mp and not {*scope}&{*mp[a]['ignore']}:
setattr(tree, i, mp[a]['to'])
n = [a] if not isinstance(a, list) else a
s = [tree.name] if tree.__class__.__name__.endswith('Def') else scope
for j in n:
if isinstance(j, ast.AST):
walk(j, mp, s)
replace_map = {
'nmp':{'to':'np', 'ignore':[]},
'x':{'to':'X', 'ignore':['main']},
'y':{'to':'Y', 'ignore':[]}
}
string = """
import numpy as nmp
y = 5
def f(x):
return nmp.sum(x) + y
x = 1
print(f(x))
"""
tree = ast.parse(string)
walk(tree, replace_map)
print(ast.unparse(tree))

输出：

import numpy as np
Y = 5
def f(X):
return np.sum(X) + Y
x = 1
print(f(x))