Python 如何通过AST重写实现代码自动注入与编译优化【技巧】

Python AST重写可在编译前修改语法树实现自动注入与优化：通过ast.NodeTransformer修改节点，compile()生成代码；支持函数级日志注入、常量折叠、无用分支剔除等，安全可靠且不依赖装饰器或eval。

Python 通过 AST（Abstract Syntax Tree）重写，可以不修改源码文本、不依赖装饰器或运行时钩子，直接在编译前干预代码结构，实现安全可靠的自动注入（如日志、权限校验）和轻量级编译优化（如常量折叠、无用代码剔除）。核心在于 ast.NodeTransformer 修改语法树，再用 compile() 生成可执行代码。

AST 注入：在函数入口/出口插入逻辑

比如为所有函数自动添加执行耗时统计。不是靠装饰器（需手动加 @timer），而是遍历函数定义节点，在其 body 开头插入计时开始语句，结尾插入结束与打印逻辑。

• 继承 ast.NodeTransformer，重写 visit_FunctionDef
• 构造 ast.Import 或 ast.Assign 节点（如 start = time.time()），用 insert(0, ...) 加到函数体开头
• 在原函数体末尾追加 ast.Expr（如 print(f"took {time.time()-start:.3f}s")）
• 注意：插入的变量名要避免污染原作用域，可用唯一前缀（如 __auto_timer_start_123）

AST 优化：识别并简化固定模式

例如把 len("hello") 直接替换成 5，或把 if False: ... 分支整个删掉——这些在解释器执行前就能完成，省去运行时开销。

• 对 ast.Call 节点判断是否是 len 调用且参数为字符串字面量，直接返回 ast.Constant(value=len(arg.s))
• 对 ast.If 节点，若测试条件是 ast.Constant 且值为 False，返回空列表（即删除该 if）；若为 True，则只保留 body 部分
• 优化后需调用 ast.fix_missing_locations() 补全行号信息，否则报错

安全注入：绕过 eval/exec，支持模块级处理

相比字符串拼接或 exec()，AST 方式天然保持语法合法性，还能跨函数/类做上下文感知注入（如只给带 @api 的函数加鉴权检查）。

• 先用 ast.parse(source) 解析整个模块，再递归扫描 ast.FunctionDef 或 ast.ClassDef
• 可通过访问父节点或 ast.get_docstring() 提取装饰器信息，实现条件注入
• 注入后用 compile(ast_tree, filename, mode='exec') 编译，再 exec() 运行——整个过程不碰原始 .py 文件

实战小技巧：调试与兼容性

AST 重写容易因节点类型或字段缺失出错，建议边写边验证。

• 用 ast.dump(tree, indent=2) 打印变换前后的树结构，对比差异
• 目标 Python 版本要匹配：3.8+ 的 ast.Constant 替代了旧版的 ast.Num/ast.Str 等，需统一处理
• 注入代码若含新语法（如 f-string），确保目标环境支持；建议注入纯表达式，避免复杂语句嵌套

基本上就这些。AST 重写不是黑魔法，它只是把“人眼改代码”的动作翻译成程序能理解的树操作。写熟了，一行注入、两行优化，比写装饰器还干净。