Python中如何让动态导入类的create函数具备类型感知？

要让create函数具备类型感知，让类型检查工具（如mypy、PyCharm）能正确推断返回实例的类型，可以通过以下步骤实现：

1. 导入类型工具并定义约束 #

在stream.py中导入Python标准库的类型注解工具，同时明确所有支持的流类与对应名称：

from typing import Type, TypeVar, Literal, Dict
from feeds.AbstractStream import AbstractStream
# 导入所有已实现的WebSocket流类
from feeds.streams import binanceusdm, okx, bybit  # 示例，根据实际类添加

# 定义TypeVar，绑定到AbstractStream的子类，确保返回类型合法
StreamSubclass = TypeVar("StreamSubclass", bound=AbstractStream)

# 用Literal限定允许传入的流名称，类型检查工具会提示非法输入
SupportedStreamNames = Literal["binanceusdm", "okx", "bybit"]

# 创建名称到流类的映射，替代getattr，提升安全性和类型推断能力
STREAM_CLASS_MAP: Dict[SupportedStreamNames, Type[AbstractStream]] = {
    "binanceusdm": binanceusdm,
    "okx": okx,
    "bybit": bybit,
}

2. 改造create函数的实现与类型注解 #

修改create函数，使用上述映射替代getattr，并添加精确的类型注解：

def create(param: SupportedStreamNames) -> Type[StreamSubclass]:
    try:
        stream_cls = STREAM_CLASS_MAP[param]
    except KeyError as e:
        raise AttributeError(f"Unknown stream: {param}") from e
    return stream_cls

效果说明 #

现在调用代码时：

s = stream.create("binanceusdm")()

类型检查工具会自动推断出s的类型是binanceusdm实例，而非Any；同时传入非法流名称时，会提前给出类型错误提示。

可选优化：自动注册子类 #

如果后续会新增大量交易所流类，手动维护映射繁琐，可以用装饰器自动注册子类：

# 初始化空的映射
STREAM_CLASS_MAP: Dict[str, Type[AbstractStream]] = {}

def register_stream(name: str):
    def decorator(cls: Type[AbstractStream]):
        STREAM_CLASS_MAP[name] = cls
        return cls
    return decorator

在每个流类文件中使用装饰器注册：

from feeds.AbstractStream import AbstractStream
from feeds.stream import register_stream

@register_stream("binanceusdm")
class binanceusdm(AbstractStream):
    # 类实现...

注：自动注册后需手动维护SupportedStreamNames的Literal类型，才能保留严格的类型检查能力。

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