目录结构

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

要替换的代码

代码在文件 utils.py 里面

def sum(a: int, b: int) -> int:c = a - breturn 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] = vreturn 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] = vreturn 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 begindef sum(a: int, b: int) -> int:c = a - breturn c-----sum code end ------- transformed code: def sum(a: int, b: int) -> int:c = a + breturn c
3 # 这是执行后的结果

最后

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