python 通过 ast 替换代码

目录结构

.
├── hello
│   ├── __init__.py
│   └── utils.py
├── main.py
├── replace_code.py

要替换的代码

代码在文件 utils.py 里面

def sum(a: int, b: int) -> int:
    c = a - b
    return c

替换代码的逻辑

替换代码逻辑在文件 replace_code.py 里

新建类，继承 ast

逻辑很简单：找到哪一行并替换，就是中间的判断可能比较多

import ast
class MyTransformer(ast.NodeTransformer):
    def visit(self, node: ast.AST):
    	# 我们要找到要替换的那一行('c = a - b')代码在哪里
        for nodei, v in enumerate(node.body):
            if isinstance(v, ast.FunctionDef):
                for i, vv in enumerate(v.body):
                    if isinstance(vv, ast.Assign) and \
                    isinstance(vv.value, ast.BinOp) and \
                    isinstance(vv.value.left, ast.Name) and \
                    vv.value.left.id == 'a' and vv.value.right.id == 'b' and isinstance(vv.value.op, ast.Sub):# 找到了那一行代码
                        # 找到后把新的代码通过 ast.parse 转换成语法树，并替换掉原来的代码
                        v.body[i] = ast.parse('c = a + b').body[0]
                        node.body[nodei] = v
                    
        return node

运行新建的类

from hello import utils
import inspect
# 获取代码内容
method_code = inspect.getsource(utils.sum)
print(f'----sum code begin\n {method_code} \n -----sum code end ----')
# 把要替换的代码转换成语法树
method_ast = ast.parse(method_code)
trans_instance = MyTransformer()
# 运行新建的类的里面的逻辑
transformed_ast = trans_instance.visit(method_ast)
# 把语法树解析成语法代码
transformed_code = ast.unparse(transformed_ast)
print('--- transformed code:', transformed_code)

最最重要的一步

一定要把代码在它所在的命名空间里执行它

exec(transformed_code, utils.__dict__)

replace_code.py 完整代码

import ast
class MyTransformer(ast.NodeTransformer):
    def visit(self, node: ast.AST):
        for nodei, v in enumerate(node.body):
            if isinstance(v, ast.FunctionDef):
                for i, vv in enumerate(v.body):
                    if isinstance(vv, ast.Assign) and \
                    isinstance(vv.value, ast.BinOp) and \
                    isinstance(vv.value.left, ast.Name) and \
                    vv.value.left.id == 'a' and vv.value.right.id == 'b' and isinstance(vv.value.op, ast.Sub):
                        print('here')
                        v.body[i] = ast.parse('c = a + b').body[0]
                        node.body[nodei] = v
                    
        return node
from hello import utils
import inspect
method_code = inspect.getsource(utils.sum)
print(f'----sum code begin\n {method_code} \n -----sum code end ----')
method_ast = ast.parse(method_code)
trans_instance = MyTransformer()
transformed_ast = trans_instance.visit(method_ast)
transformed_code = ast.unparse(transformed_ast)
print('--- transformed code:', transformed_code)
exec(transformed_code, utils.__dict__)

main.py 里面的代码

导入的顺序一定不要错了，不然也不会生效的

# 引入的顺序一定不要错了，不然也不会生效的
import replace_code
from hello.utils import sum
if __name__ == '__main__':
    print(sum(2,1))

执行结果

----sum code begin
 def sum(a: int, b: int) -> int:
    c = a - b
    return c
 
 -----sum code end ----

--- transformed code: def sum(a: int, b: int) -> int:
    c = a + b
    return c
3 # 这是执行后的结果

最后

逻辑代码不一定要在 visit() 里面写，如果你要修改的是 for 循环里面的代码 你可以重写 visit_For()方法， 在里面来实现自己的逻辑