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:

Answer 2

Pylance is a static analysis tool, so with particularly complex cases like this, it can be extremely challenging to get it to fully understand your code.

The reason that your current code doesn't work is because Pylance has no way of determining whether the value of factory_instance.version has changed between instantiating the Factory and using it to instantiate sub-classes, since the values of object attributes aren't persisted between function calls during the analysis.

This means that it is impossible to make Pylance (or other static analysis tools) understand the code without you manually adding significant amounts of supporting code, even if you use @overload decorators.

Fortunately, as shown in this answer, it is very possible to write this supporting code, although it can be very tedious to do so.

In your particular case, you need to trick Pylance into thinking that instantiating your class with different variants of Version will give completely different subclasses.

Note that in reality, you are still returning instances of Factory, but the @overload definitions are enough to trick static analysis tools into thinking that they are different.

class Factory:
    @overload
    def __new__(cls, version: Literal[Version.Version1]) -> '__Factory_V1':
        ...

    @overload
    def __new__(cls, version: Literal[Version.Version2]) -> '__Factory_V2':
        ...

    @overload
    def __new__(cls, version: Literal[Version.Version3]) -> '__Factory_V3':
        ...

    def __new__(cls, version: Version) -> Self:
        # This is just the default implementation
        # We're actually instantiating the object normally
        return super().__new__(cls)

    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)

From there, getting proper editor support is just a matter of writing subclasses which specifically declare the return types. Notice that we are preventing accesses to the properties by raising a NotImplementedError to ensure that people can't actually use our methods.

class __Factory_V1(Factory):
    @property
    def Machine1Config(self) -> type[Machine1Config_1]:
        raise NotImplementedError("This subclass should never be instantiated")

    @property
    def Machine2Config(self) -> type[Machine2Config_1]:
        raise NotImplementedError("This subclass should never be instantiated")


class __Factory_V2(Factory):
    @property
    def Machine1Config(self) -> type[Machine1Config_2]:
        raise NotImplementedError("This subclass should never be instantiated")

    @property
    def Machine2Config(self) -> type[Machine2Config_2]:
        raise NotImplementedError("This subclass should never be instantiated")


class __Factory_V3(Factory):
    @property
    def Machine1Config(self) -> type[Machine1Config_2]:
        raise NotImplementedError("This subclass should never be instantiated")

    @property
    def Machine2Config(self) -> type[Machine2Config_2]:
        raise NotImplementedError("This subclass should never be instantiated")

This will result in your instantiations having the correct type type definitions.

factory_v1_instance = Factory(Version.Version1)
machine1_config_instance = factory_v1_instance.Machine1Config("petrol", 100)

reveal_type(machine1_config_instance)
# Pylance: Type of "machine1_config_instance" is "Machine1Config_1"

factory_v2_instance = Factory(Version.Version2)
machine2_config_instance = factory_instance.Machine2Config("petrol", 100)

reveal_type(machine2_config_instance)
# Pylance: Type of "machine2_config_instance" is "Machine2Config_2"

Note that this doesn't appear to work correctly in Mypy. Since you specifically asked for Pylance, I haven't bothered investigating why, but I have managed to get things like this working with it before (see the answer I linked above), so it is probably possible with some fiddling. My best guess is that it doesn't check @overload decorations on the __new__ method.

reveal_type(machine1_config_instance)
# Mypy: Revealed type is "Any"

As you've likely observed, this approach is very tedious, and requires the types to be declared manually. I wouldn't recommend it unless getting top-notch editor support for this type is absolutely essential, since it makes the code nightmarish to maintain due to the need to update definitions in many places. However, it certainly does work, at least for Pylance.

Answer 3

Because you know what you are getting, you can actually tell pylance via the cast (from typing) function.

You last lines becomes:

machine1_config_instance = cast(Type[Machine1Config_1], factory_instance.Machine1Config())() # completion works correctly here
machine2_config_instance = cast(Type[Machine2Config_1], factory_instance.Machine2Config())() # completion works correctly here

You are basically manually saying that you know the return type, that it is a type, not object, so it is callable.

Is it clunky? Yes. Does it work? Yes as well. Is it worth it? Not sure. Is there a better way? Possibly, but as Pylance (Pyright in this case) is a static tool, you will either need to feed it something manually, or possibly do some redesign.