如何根据__init__参数值为类属性jobs添加对应类型提示

要让静态类型检查器（比如Pylance）根据input_type的具体值推断jobs的精确类型，你可以通过重载__init__方法或者使用泛型类来实现，以下是两种可行方案：

方案一：重载__init__方法 #

直接为__init__添加多个@overload签名，明确不同input_type对应的jobs类型：

from typing import Literal, overload

class VideoClass:
    def video_method(self):
        pass

class ImageClass:
    def image_method(self):
        pass

class TextClass:
    def text_method(self):
        pass

InputType = Literal["TEXT", "IMAGE", "VIDEO"]

class MyClass:
    @overload
    def __init__(self, input_type: Literal["VIDEO"]) -> None:
        ...
    @overload
    def __init__(self, input_type: Literal["IMAGE"]) -> None:
        ...
    @overload
    def __init__(self, input_type: Literal["TEXT"]) -> None:
        ...
    def __init__(self, input_type: InputType) -> None:
        if input_type == "VIDEO":
            self.jobs = VideoClass()
        elif input_type == "IMAGE":
            self.jobs = ImageClass()
        elif input_type == "TEXT":
            self.jobs = TextClass()

# 类型检查示例
video_instance = MyClass("VIDEO")
video_instance.jobs.video_method()  # Pylance会识别jobs为VideoClass，无报错

image_instance = MyClass("IMAGE")
image_instance.jobs.image_method()  # 同理识别为ImageClass

方案二：使用泛型类 #

通过泛型绑定input_type和jobs的类型映射，让类型检查器关联两者：

from typing import Literal, Generic, TypeVar

class VideoClass:
    def video_method(self):
        pass

class ImageClass:
    def image_method(self):
        pass

class TextClass:
    def text_method(self):
        pass

# 定义类型变量，绑定输入类型和对应的类类型
T = TypeVar("T", VideoClass, ImageClass, TextClass)
InputTypeMap = {
    "TEXT": TextClass,
    "IMAGE": ImageClass,
    "VIDEO": VideoClass,
}

class MyClass(Generic[T]):
    def __init__(self, input_type: Literal["TEXT", "IMAGE", "VIDEO"]) -> None:
        self.jobs: T = InputTypeMap[input_type]()

# 类型检查示例
video_instance: MyClass[VideoClass] = MyClass("VIDEO")
video_instance.jobs.video_method()

image_instance: MyClass[ImageClass] = MyClass("IMAGE")
image_instance.jobs.image_method()

注意事项 #

• 方案一不需要显式指定泛型参数，更贴合原有代码的使用习惯，Pylance能自动根据传入的input_type推断jobs类型。
• 方案二需要在实例化时显式指定泛型类型，适合需要更明确类型约束的场景。

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