如何优化工厂模式代码实现Pylance动态类参数精准自动补全

问题场景还原 #

先看一下典型的问题代码（你的工厂模式结构大概类似这样）：

from typing import Union, Type
from enum import Enum

class Version(Enum):
    V1 = "v1"
    V2 = "v2"

class Machine1Config:
    def __init__(self, param_a: str, param_b: int):
        self.param_a = param_a
        self.param_b = param_b

class Machine2Config:
    def __init__(self, param_c: bool, param_d: float):
        self.param_c = param_c
        self.param_d = param_d

class ConfigFactory:
    @staticmethod
    def create_config(version: Version) -> Union[Machine1Config, Machine2Config]:
        config_cls: Type[Union[Machine1Config, Machine2Config]]
        if version == Version.V1:
            config_cls = Machine1Config
        else:
            config_cls = Machine2Config
        return config_cls()  # 这里补全会同时显示两个类的所有参数，混乱且不准确

优化方案：用@overload实现精准类型推断 #

通过为工厂方法添加重载签名，让Pylance/Pyright能根据传入的Version枚举值，精准推断返回的配置类类型，从而提供对应类的专属参数补全：

from typing import overload, TypeVar
from enum import Enum

class Version(Enum):
    V1 = "v1"
    V2 = "v2"

class Machine1Config:
    def __init__(self, param_a: str, param_b: int):
        self.param_a = param_a
        self.param_b = param_b

class Machine2Config:
    def __init__(self, param_c: bool, param_d: float):
        self.param_c = param_c
        self.param_d = param_d

# 定义类型变量，用于绑定重载时的具体返回类型
T = TypeVar("T", Machine1Config, Machine2Config)

class ConfigFactory:
    # 为每个Version枚举值单独定义重载签名
    @overload
    @staticmethod
    def create_config(version: Version.V1) -> Machine1Config: ...

    @overload
    @staticmethod
    def create_config(version: Version.V2) -> Machine2Config: ...

    # 实际运行时的实现方法
    @staticmethod
    def create_config(version: Version) -> Machine1Config | Machine2Config:
        if version == Version.V1:
            return Machine1Config()
        else:
            return Machine2Config()

进阶：支持直接传递构造参数 #

如果需要在调用工厂方法时直接传入配置类的构造参数，可进一步重载签名：

from typing import overload, TypeVar
from enum import Enum

class Version(Enum):
    V1 = "v1"
    V2 = "v2"

class Machine1Config:
    def __init__(self, param_a: str, param_b: int):
        self.param_a = param_a
        self.param_b = param_b

class Machine2Config:
    def __init__(self, param_c: bool, param_d: float):
        self.param_c = param_c
        self.param_d = param_d

T = TypeVar("T", Machine1Config, Machine2Config)

class ConfigFactory:
    @overload
    @staticmethod
    def create_config(version: Version.V1, param_a: str, param_b: int) -> Machine1Config: ...

    @overload
    @staticmethod
    def create_config(version: Version.V2, param_c: bool, param_d: float) -> Machine2Config: ...

    @staticmethod
    def create_config(version: Version, *args, **kwargs) -> Machine1Config | Machine2Config:
        if version == Version.V1:
            return Machine1Config(*args, **kwargs)
        else:
            return Machine2Config(*args, **kwargs)

关键说明 #

1. @overload签名仅用于类型检查器（Pylance/Pyright）识别，不需要实现逻辑，用...占位即可。
2. 类型检查器会根据传入的具体Version值，自动匹配对应的重载签名，从而给出精准的参数补全提示。
3. 若使用Python 3.9及以下版本，将Machine1Config | Machine2Config替换为Union[Machine1Config, Machine2Config]即可。
4. 确保VS Code中Pylance的类型检查模式开启（设置项python.analysis.typeCheckingMode设为basic或strict）。

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