【Bug已解决】Add type hints to public API functions 解决方案-尧图网络科技

【Bug已解决】Add type hints to public API functions 解决方案

一、现象长什么样

这是一类「非崩溃但严重拖累工程健康度」的问题：一个库的公开 API（public API）函数没有类型注解（type hints），导致：

IDE 无法自动补全参数名与返回类型，使用者只能翻源码；
静态检查工具（mypy / pyright / pyright in VS Code）对调用处全部放行，隐藏大量拼错参数、传错类型的隐患；
重构时「改了函数签名却没人发现调用方不匹配」，直到运行时才炸；
文档生成工具（如 Sphinx autodoc）拿不到类型信息，文档里参数类型全是Any。

在 DeepSpeed 这类大型 C 扩展 + Python 封装的项目里尤其明显：很多deepspeed.xxx顶层函数、初始化入口、工具函数完全无注解，新贡献者贡献代码时常踩「参数顺序记错」「返回值类型不清」的坑。维护者于是开了 issue 请求「给公开 API 加类型注解」。

本期讲清楚：为什么缺类型提示是真实 Bug、如何在不破坏兼容的前提下系统地补上类型提示，以及三层工程化做法。

二、背景

2.1 类型提示的价值

Python 3.5+ 引入的 PEP 484 类型提示，本质是「可选的静态契约」：

def add(a: int, b: int) -> int: return a + b

它不影响运行时行为，但让mypy能在不跑代码的情况下发现add("1", 2)这类错误。对公开库来说，类型提示是 API 契约的一部分——使用者靠它理解「传什么、得到什么」。

2.2 为什么大型库容易缺类型提示

历史包袱：很多函数写于类型提示普及前；
C 扩展封装：底层是 C/C++，Python 层只是薄封装，类型推断困难；
动态返回：同一函数根据配置返回不同类型（如deepspeed.initialize返回(model, optimizer, ...)元组，或不同 engine 子类），注解复杂；
担心破坏：from __future__ import annotations引入前，注解会在模块加载时求值，引用尚未定义的类会报NameError。

三、根因

3.1 缺注解 -> 调用方无保护

没有注解时，下面这种错误谁也拦不住：

# 库函数(无注解) def initialize(model, config, parameters=None): ... # 用户调用 engine = initialize(config=my_config, model=my_model, lr=1e-3) # lr 不是 initialize 的参数, 但运行时才因意外 kwarg 报错

如果initialize有注解且用 mypy 检查，lr=1e-3会在静态阶段被指出。

3.2 动态返回类型造成「调用方只能 guess」

deepspeed.initialize返回(model, optimizer, _, lr_scheduler)这种异构元组，无注解时使用者不知道第 3 个元素是什么、能不能忽略，只能看例子照抄，极易出错。

3.3 Python 3.9 及以下的前向引用坑

在 Python 3.9 里直接写def f() -> MyCls:而MyCls在文件后面才定义，会NameError。这也劝退了很多贡献者去加注解——其实有标准解法（见下文）。

3.4 一句话根因

公开 API 缺类型提示，使调用方失去静态契约保护，参数拼错、类型传错只能等到运行时暴露，同时 IDE 补全与文档生成失效；而「前向引用、动态返回、C 扩展封装」等技术顾虑又让维护者迟迟不愿补——最终形成工程健康度负债。

四、最小可运行复现

下面演示「无注解 -> mypy 放行错误调用」与「加注解 -> 静态拦截」的对比：

# ---- 无注解版本 ---- def divide(a, b): return a / b divide("10", 2) # mypy 不报错(因为无注解), 运行时才 TypeError # ---- 加注解版本 ---- def divide(a: float, b: float) -> float: return a / b # divide("10", 2) # 取消注释后 mypy 会报: Argument 1 to "divide" has incompatible type "str"; expected "float"

用 mypy 跑：

pip install mypy mypy demo.py

无注解版：Success: no issues found（错误被放过）。有注解版（含错误调用）：error: Argument 1 ... incompatible type "str"; expected "float"。

这就是类型提示把「运行时崩溃」提前到「提交前」的价值。

五、解决方案（第一层：最小直接修复）

给公开函数逐一补注解。对大多数纯 Python 封装函数，直接标注即可：

from typing import Optional, Dict, Any, Tuple def get_argument( name: str, default: Optional[Any] = None, dtype: type = str, ) -> Any: ... def initialize( model: "torch.nn.Module", config: Dict[str, Any], parameters: Optional[list] = None, ) -> Tuple[Any, Any, Any, Any]: ...

5.1 解决前向引用（Python 3.9 兼容）

DeepSpeed 仍需支持 Python 3.9，不能直接用model: torch.nn.Module当torch在文件顶部未 import 完成时求值。用字符串注解或from __future__ import annotations：

from __future__ import annotations # 让所有注解变成字符串, 延迟求值, 兼容 3.9 from typing import Optional class Engine: ... def build(opts: Optional[Engine]) -> Engine: # 即使 Engine 后定义也 OK ...

from __future__ import annotations是 Python 3.7+ 可用、3.9 完全支持的写法，是给老项目补注解的最优解。

六、解决方案（第二层：结构性 / 抽象改进）

第一层是「手写注解」，但更系统的是引入 Protocol / 类型别名，统一管理复杂返回，并把公开 API 收敛到少量入口。

6.1 用 Protocol 描述复杂对象

deepspeed.initialize返回的 engine 有forward、backward、step等方法，可用Protocol描述，供调用方获得补全：

from __future__ import annotations from typing import Protocol, Any, Tuple, Optional class DeepSpeedEngineProtocol(Protocol): def forward(self, *args, **kwargs) -> Any: ... def backward(self, loss: Any) -> None: ... def step(self) -> None: ... def initialize( model: Any, config: dict, parameters: Optional[list] = None, ) -> Tuple[DeepSpeedEngineProtocol, Any, Any, Any]: ...

6.2 pydantic 配置模型替代裸 dict

DeepSpeed 配置是嵌套 dict，无类型导致config["zero_optimization"]["stage"]拼错无提示。用 pydantic 模型承载：

from __future__ import annotations from pydantic import BaseModel, Field class ZeroConfig(BaseModel): stage: int = Field(ge=0, le=3) offload_param: dict = Field(default_factory=dict) class DSConfig(BaseModel): zero_optimization: ZeroConfig = Field(default_factory=ZeroConfig) fp16: dict = Field(default_factory=dict) # 调用方拿到 DSConfig, IDE 能补全 .zero_optimization.stage

6.3 用 pyright 的`py.typed`标记

纯 Python 库要在包根放一个空的py.typed文件，类型检查器才会把你的注解当作「对外契约」：

touch deepspeed/py.typed

并在pyproject.toml里把它纳入打包package-data。

七、解决方案（第三层：断言 / CI 守护）

把「公开 API 必须有注解」变成 CI 不变量。

7.1 mypy 严格模式接入 CI

# pyproject.toml [tool.mypy] python_version = "3.9" disallow_untyped_defs = true # 禁止无注解函数 disallow_incomplete_defs = true # 禁止部分注解 warn_return_any = true ignore_missing_imports = true

# CI jobs: type-check: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - run: pip install mypy - run: mypy deepspeed/

disallow_untyped_defs = true会让任何新增的无注解公开函数直接 CI 失败，从而保证「补注解」不退化。

7.2 用脚本统计未注解的公开函数

import ast, pathlib def count_untyped_public(path: str) -> list: tree = ast.parse(pathlib.Path(path).read_text()) issues = [] for node in ast.walk(tree): if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): if node.name.startswith("_"): continue # 私有跳过 args = node.args all_args = [*args.args, *args.kwonlyargs] untyped = [a.arg for a in all_args if a.annotation is None] if node.returns is None: untyped.append("-> return") if untyped: issues.append((node.lineno, node.name, untyped)) return issues if __name__ == "__main__": for ln, name, missing in count_untyped_public("deepspeed/__init__.py"): print(f"L{ln} {name} 缺注解: {missing}")

三层叠加：直接补注解（含from __future__ import annotations）+ 结构改（Protocol/pydantic/py.typed）+ 守护（mypy 严格 + 统计脚本 CI），类型提示从「负债」变成「契约」。

八、补充：注解不影响运行时，但要注意 eval 时机

两个常见坑：

注解里引用未导入的名字会NameError（Python 3.9 无from __future__ import annotations时）。解决：加from __future__ import annotations或改成字符串注解"MyClass"。
from __future__ import annotations下，运行时typing.get_type_hints仍会求值字符串——若字符串引用的类在模块外且未 import，会在反射时失败。确保相关类可被 import。

另外，类型提示只是「提示」，运行时不做强制校验。若需要运行时校验（比如配置 dict 的字段），应另用 pydantic /dataclasses+ 显式assert。

九、排查清单

当团队决定「给公开 API 加类型提示」时：

先加from __future__ import annotations，规避 Python 3.9 前向引用NameError。
从顶层入口函数开始（如initialize、get_argument），逐步向内。
复杂返回用Protocol/TypeVar描述，给调用方补全。
配置类用 pydantic / dataclass承载，替代裸dict。
放py.typed标记并打包，让外部项目能用你的注解。
CI 开disallow_untyped_defs，禁止新增无注解函数。
写 AST 统计脚本，定期列出剩余未注解的公开函数，量化进度。
注意注解 eval 时机，避免get_type_hints因未 import 而失败。

十、小结

「公开 API 缺类型提示」看似不是崩溃型 bug，却真实损害工程质量：调用方失去静态契约、IDE 补全失效、文档类型缺失、重构隐患只能运行时暴露。根因在于历史包袱、前向引用顾虑、动态返回与 C 扩展封装让维护者迟迟未补。

修复分三层：第一层，用from __future__ import annotations逐个给公开函数加注解（兼容 Python 3.9）；第二层，用Protocol描述复杂返回、pydantic承载配置、py.typed标记对外契约；第三层，mypy 严格模式 + AST 统计脚本接入 CI，让「无注解公开函数」成为不可合入的失败。类型提示一旦成为工程习惯，大量「运行时才发现的拼错参数」都会被提前到「提交前」拦截。