如何基于ast.NodeTransformer自动向Python项目所有模块注入代码并运行项目

既然你已经能通过ast.NodeTransformer处理单个模块，那核心问题就是在模块被导入前拦截并修改AST，而且不用第三方库——其实Python标准库的importlib和ast模块就能搞定，完全不需要astor这类工具。下面是具体的实现思路和代码示例：

核心思路：拦截Python的导入流程 #

Python的导入系统是可扩展的，我们可以通过自定义导入钩子（import hooks），在模块被编译成字节码前修改它的AST，然后再执行修改后的代码。整个过程完全在内存中完成，不会修改原项目文件，也不需要把AST转回源码。

具体实现步骤 #

1. 定义你的AST修改器 #

先把你已经写好的NodeTransformer拿过来，比如这个简单的示例（在每个函数开头注入打印语句）：

import ast

class MyCodeInjector(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        # 注入一行打印代码，注意要解析成AST节点
        inject_stmt = ast.parse("print(f'Injected: {__name__}.{node.name}')").body[0]
        node.body.insert(0, inject_stmt)
        return self.generic_visit(node)

2. 实现自定义模块加载器 #

这个加载器负责读取模块源码、修改AST、编译并执行修改后的代码：

import importlib.abc
import importlib.util
import sys
import os

class ModifiedLoader(importlib.abc.SourceLoader):
    def __init__(self, module_path):
        self.module_path = module_path

    def get_source(self, fullname):
        # 读取模块源码
        with open(self.module_path, 'r', encoding='utf-8') as f:
            return f.read()

    def get_filename(self, fullname):
        return self.module_path

    def exec_module(self, module):
        # 1. 解析源码为AST
        source = self.get_source(module.__name__)
        tree = ast.parse(source)

        # 2. 用自定义修改器处理AST
        transformer = MyCodeInjector()
        modified_tree = transformer.visit(tree)
        # 修复新节点的位置信息，避免编译报错
        ast.fix_missing_locations(modified_tree)

        # 3. 编译修改后的AST为字节码
        bytecode = compile(modified_tree, filename=self.module_path, mode='exec')

        # 4. 执行字节码，填充模块对象
        exec(bytecode, module.__dict__)

3. 实现自定义模块查找器 #

这个查找器负责识别项目内的模块，并使用上面的自定义加载器来加载它们：

class ProjectFinder(importlib.abc.MetaPathFinder):
    def __init__(self, project_root):
        self.project_root = os.path.abspath(project_root)

    def find_spec(self, fullname, path, target=None):
        # 把模块名转换成文件路径（比如"my_package.module" → "my_package/module.py"）
        module_parts = fullname.split('.')
        module_path = os.path.join(self.project_root, *module_parts) + '.py'

        # 只处理项目根目录下存在的模块
        if os.path.exists(module_path) and os.path.abspath(module_path).startswith(self.project_root):
            # 创建模块规格，指定用自定义加载器
            spec = importlib.util.spec_from_loader(
                fullname,
                ModifiedLoader(module_path),
                origin=module_path
            )
            return spec
        # 非项目内的模块，交给Python默认导入机制处理
        return None

4. 编写项目入口脚本 #

在导入任何项目模块前，先注册自定义的导入钩子，然后再导入主模块运行：

if __name__ == '__main__':
    # 设置你的项目根目录（比如当前脚本所在的目录）
    PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))

    # 把自定义查找器添加到导入路径最前面，确保优先被调用
    sys.meta_path.insert(0, ProjectFinder(PROJECT_ROOT))

    # 现在导入并运行主模块，所有被它导入的项目内模块都会被修改
    import my_main_module
    my_main_module.main()

关键注意事项 #

• 钩子要先注册：一定要在导入任何项目模块前设置sys.meta_path，否则已经导入的模块不会被处理。
• 处理包结构：如果你的项目有子包（带__init__.py的目录），可以扩展find_spec的逻辑，比如检查目录下的__init__.py文件。
• 避免修改外部模块：查找器里的路径检查确保只处理你项目内的模块，不会修改标准库或第三方依赖。
• 位置信息修复：修改AST后必须调用ast.fix_missing_locations，否则新添加的节点没有位置信息，编译时会报错。

这个方案完全基于Python标准库，没有第三方依赖，而且是在内存中临时修改模块，不会影响你的原代码文件，完美匹配你的需求。

内容的提问来源于stack exchange，提问作者Gal Shahar