Pylance not working autocomplete for dynamically instantiated classes

Question

from typing import Literal, overload, TypeVar, Generic, Type
import enum
import abc
import typing

class Version(enum.Enum):
    Version1 = 1
    Version2 = 2
    Version3 = 3


import abc
from typing import Type


class Machine1BaseConfig:
    @abc.abstractmethod
    def __init__(self, *args, **kwargs) -> None:
        pass

class Machine1Config_1(Machine1BaseConfig):
    def __init__(self, fueltype, speed) -> None:
        self.fueltype = fueltype
        self.speed = speed

class Machine1Config_2(Machine1BaseConfig):
    def __init__(self, speed, weight) -> None:
        self.speed = speed
        self.weight = weight

class Machine1FacadeConfig:
    @classmethod
    def get_version(cls, version: Version) -> Type[typing.Union[Machine1Config_1, Machine1Config_2]]:
        config_map = {
            Version.Version1: Machine1Config_1,
            Version.Version2: Machine1Config_2,
            Version.Version3: Machine1Config_2,
        }
        return config_map[version]
        

class Machine2BaseConfig:
    @abc.abstractmethod
    def __init__(self, *args, **kwargs) -> None:
        pass

class Machine2Config_1(Machine2BaseConfig):
    def __init__(self, gridsize) -> None:
        self.gridsize = gridsize

class Machine2Config_2(Machine2BaseConfig):
    def __init__(self, loadtype, duration) -> None:
        self.loadtype = loadtype
        self.duration = duration

class Machine2FacadeConfig:
    @classmethod
    def get_version(cls, version: Version) -> Type[typing.Union[Machine2Config_1, Machine2Config_2]]:
        config_map = {
            Version.Version1: Machine2Config_1,
            Version.Version2: Machine2Config_1,
            Version.Version3: Machine2Config_2,
        }
        return config_map[version]


class Factory:
    def __init__(self, version: Version) -> None:
        self.version = version

    @property
    def Machine1Config(self):
        return Machine1FacadeConfig.get_version(self.version)

    @property
    def Machine2Config(self):
        return Machine2FacadeConfig.get_version(self.version)


factory_instance = Factory(Version.Version1)
machine1_config_instance = factory_instance.Machine1Config()
machine2_config_instance = factory_instance.Machine2Config()

In the provided Python code, the Factory class is used to instantiate configuration objects for two different types of machines (Machine1 and Machine2) based on a specified version. The problem is when using Pylance/Pyright with Visual Studio Code, I'm experiencing issues with autocomplete not correctly suggesting parameters for dynamically instantiated classes (Machine1Config and Machine2Config) in a factory design pattern. How can I improve my code to enable more accurate and helpful autocompletion suggestions by Pylance for these dynamically determined types?

I have thought that this should somehow work with @overload decorater but I can't wrap my head around it how to quite implement it.

Furthermore currently with the type hint Type[typing.Union[Machine1Config_1, Machine1Config_2]] Pylance suggests all key word arguments of Machine1Config_1 and Machine1Config_2, so fueltype, speed, weight. If I leave this type hint away there is no autocompletion at all.

Answer 1

Looking at the factory, there is no way to tell which of Type[typing.Union[Machine2Config_1, Machine2Config_2]] will be returned when calling Machine1FacadeConfig.get_version(self.version) in isolation.

As the facade and the factory are extremely coupled anyways, I would suggest combining these into a single utility, where the types for version and configs can be more tightly coupled.

You can declare a generic class for factory and provide a helper function which returns an instance of that factory where the version and config types have been bound together. The helper function would be overloaded for the different combinations.

_Config1 = TypeVar("_Config1", Machine1Config_1, Machine1Config_2)
_Config2 = TypeVar("_Config2", Machine2Config_1, Machine2Config_2)


class _Factory(Generic[_Config1, _Config2]):
    def __init__(self, config1: Type[_Config1], config2: Type[_Config2]):
        self._config1 = config1
        self._config2 = config2

    @property
    def Machine1Config(self) -> Type[_Config1]:
        return self._config1

    @property
    def Machine2Config(self) -> Type[_Config2]:
        return self._config2


@overload
def Factory(version: Literal[Version.Version1]) -> _Factory[Machine1Config_1, Machine2Config_1]:
    ...

@overload
def Factory(version: Literal[Version.Version2]) -> _Factory[Machine1Config_1, Machine2Config_2]:
    ...

@overload
def Factory(version: Literal[Version.Version3]) -> _Factory[Machine1Config_2, Machine2Config_2]:
    ...


def Factory(version: Version) -> _Factory:
    config_map1 = {
        Version.Version1: Machine1Config_1,
        Version.Version2: Machine1Config_1,
        Version.Version3: Machine1Config_2,
    }
    config_map2 = {
        Version.Version1: Machine2Config_1,
        Version.Version2: Machine2Config_2,
        Version.Version3: Machine2Config_2,
    }
    return _Factory(config_map1[version], config_map2[version])


factory_instance = Factory(Version.Version1)
machine1_config_instance = factory_instance.Machine1Config
machine2_config_instance = factory_instance.Machine2Config

Example screenshot from vscode: