I have some Python code as a string,
string = """
import numpy as nmp
y = 5
def f(x):
return nmp.sum(x) + y
x = 1
print(f(x))
"""
Using Python, I would like to rename
the import to nmp to np,
the function argument x to X (but not the x = 1), and
y to Y.
To this end, I must be able to identify every usage of the respective variables. I suppose I'd have to use one of ast or libcst, but I'm not really sure.
Any hints?
The code below traverses the source code as an ast.AST object, updating names based on a provided mapping. This mapping lists the new names for the objects to be mutated, and also specifies any scopes in which the replacement should be prohibited:
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))
Output:
import numpy as np
Y = 5
def f(X):
return np.sum(X) + Y
x = 1
print(f(x))
