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 的函数加鉴权检查）。

标签： python node 作用域