为无法直接导入的C扩展模块A.B生成pyi存根文件

问题背景 #

你有一个基于Boost编写的C扩展包A，其中A.B并非标准子模块，而是作为A的属性暴露。直接import A.B会触发ModuleNotFoundError，仅能通过from A import B或import A后访问A.B获取模块。同时B的子模块在sys.modules中以扁平路径（如B.x.y.z而非A.B.x.y.z）存在，导致pyright --createstub或mypy stubgen因无法正常导入A.B而生成存根失败。

可行解决方案 #

方案1：修改mypy stubgen的导入逻辑 #

stubgen默认通过importlib.import_module("A.B")导入目标模块，我们可以修改其导入逻辑以适配这种特殊模块：

1. 找到mypy源码中stubgen.py文件的import_module函数
2. 添加针对A.B及其子模块的特殊处理：

   def import_module(name: str) -> ModuleType:
       import sys
       import importlib
       # 处理A.B本身
       if name == "A.B":
           import A
           mod = A.B
           sys.modules[name] = mod
           return mod
       # 处理A.B的子模块（如A.B.x.y.z）
       elif name.startswith("A.B."):
           import A
           # 拆分模块路径，跳过前两部分"A"和"B"
           subparts = name.split(".")[2:]
           current_mod = A.B
           for part in subparts:
               current_mod = getattr(current_mod, part)
           # 将模块注册到sys.modules对应路径
           sys.modules[name] = current_mod
           return current_mod
       # 其他模块用默认导入逻辑
       return importlib.import_module(name)
   

3. 编译修改后的mypy，或直接运行修改后的stubgen.py生成存根：

   python path/to/modified/stubgen.py -o stubs A.B
   

方案2：使用自定义导入钩子（无需修改stubgen源码） #

通过Python的导入钩子机制，在stubgen运行前拦截A.B的导入请求，手动映射到正确的模块：

1. 创建预加载脚本preload_stubgen.py：

   import sys
   import importlib.abc
   import importlib.util
   import subprocess
   
   class ABModuleImporter(importlib.abc.MetaPathFinder):
       def find_spec(self, fullname, path, target=None):
           # 处理A.B及子模块的导入请求
           if fullname == "A.B":
               import A
               mod = A.B
               sys.modules[fullname] = mod
           elif fullname.startswith("A.B."):
               # 映射到sys.modules中的扁平路径模块（如A.B.x.y → B.x.y）
               flat_mod_name = fullname.replace("A.B.", "B.", 1)
               # 确保A已导入，子模块已加载到sys.modules
               import A
               mod = sys.modules.get(flat_mod_name)
               if mod:
                   sys.modules[fullname] = mod
           # 返回自定义spec让Python使用我们注册的模块
           if fullname in sys.modules:
               return importlib.util.spec_from_loader(fullname, None, origin="boost-ext-module")
           return None
   
       def create_module(self, spec):
           return sys.modules[spec.name]
   
       def exec_module(self, module):
           pass
   
   # 注册导入钩子
   sys.meta_path.insert(0, ABModuleImporter())
   
   # 启动stubgen生成存根
   subprocess.run(
       ["stubgen", "-o", "./stubs", "A.B"],
       env=dict(sys.environ, PYTHONPATH=".".join(sys.path))
   )
   

2. 运行该脚本：

   python preload_stubgen.py
   

方案3：预加载模块并注册到sys.modules #

如果不想写导入钩子，可以手动预加载模块并将A.B及其子模块注册到对应路径，再启动stubgen：

1. 创建脚本preload_and_stub.py：

   import sys
   import importlib
   import subprocess
   
   # 预加载A，确保B及其子模块被加载到sys.modules
   import A
   
   # 手动注册A.B到sys.modules
   sys.modules["A.B"] = A.B
   
   # 递归注册所有子模块（示例，可根据实际结构调整）
   def register_submodules(parent_mod, parent_name):
       for attr_name in dir(parent_mod):
           attr = getattr(parent_mod, attr_name)
           if isinstance(attr, type(sys)) and attr.__name__ != parent_name:
               mod_name = f"{parent_name}.{attr_name}"
               if mod_name not in sys.modules:
                   sys.modules[mod_name] = attr
                   register_submodules(attr, mod_name)
   
   register_submodules(A.B, "A.B")
   
   # 运行stubgen
   subprocess.run(["stubgen", "-o", "stubs", "A.B"])
   

2. 执行脚本生成存根。

注意事项 #

• 如果B的子模块结构复杂，递归注册时需避免循环导入或非模块属性被误注册
• 修改mypy源码时，建议基于对应版本的分支修改，避免版本兼容问题
• 生成存根后，可手动调整存根文件中的模块结构，确保类型提示准确

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