Search code examples
pythonvisual-studio-codepython-typingpylance

Is there a way for static analyzers (e.g. VS Code Intellisense) to recognize classes created through class factories?


I have the following code.

from abc import ABC, abstractmethod
from typing import *


class ExampleClass(ABC):
    @abstractmethod
    def get(self):
        pass


def class_factory(s: str) -> Type[ExampleClass]:
    class SpecificClass(ExampleClass):
        def __init__(self):
            pass

        def get(self):
            return s

    return SpecificClass


ASubclass = class_factory("A") 
BSubclass = class_factory("B")


def f(x: ASubclass):
    pass

In VS Code, in f, the type annotation ASubclass has a red underline with the message

Variable not allowed in type expression Pylance(reportGeneralTypeIssues)

When I hover over the definition of ASubclass, VS Code is able to recognize that it is a Type[ExampleClass], so I would expect at the very least when coding f, typing something like x. would give .get() as a suggestion, but when I hover over x, VS Code tells me x: Unknown.

What I want is the behavior of VS Code (or any static analyzer) with this code, without having to potentially write out tens or hundreds of potentially similar classes.

from abc import ABC, abstractmethod
from typing import *


class ExampleClass(ABC):
    @abstractmethod
    def get(self):
        pass


class ASubclass(ExampleClass):
    def get(self):
        return "A"
    

class BSubclass(ExampleClass):
    def get(self):
        return "B"


def f(x: ASubclass):
    pass

How can this be achieved?


Solution

  • There's a couple of ways to go on this - an important thing to remember is that types can't borrow information from runtime behaviour. The reason your above example doesn't work is that it's trying to use a runtime type as a type alias.

    Solution 1: Use the base class

    The simplest way to go, is probably to use ExampleClass as the type in f - it doesn't tell the user that the get method returns "A", but otherwise it does the job:

    from abc import ABC, abstractmethod
    from typing import *
    
    
    class ExampleClass(ABC):
        @abstractmethod
        def get(self) -> str:
            ...
    
    
    def class_factory(s: str) -> type[ExampleClass]:
        class SpecificClass(ExampleClass):
            def __init__(self):
                pass
    
            def get(self) -> str:
                return s
    
        return SpecificClass
    
    
    ASubclass = class_factory("A") 
    BSubclass = class_factory("B")
    
    
    def f(x: ExampleClass):
        pass
    
    f(ASubclass())
    

    Solution 2: Use TypeVar

    So for whatever reason, let's suppose you really need to know the type on the get return is more specific than just string. TypeVar can come to our aid here. We can use it to specify the return type of get varies depending on the inheritance structure using Generic, and this allows us to specify a class where the return type of get is specifically the string "A":

    from abc import ABC, abstractmethod
    from typing import Generic, Literal, LiteralString, TypeVar
    
    
    GetReturn = TypeVar("GetReturn", bound=LiteralString)
    class ExampleClass(Generic[GetReturn], ABC):
        @abstractmethod
        def get(self) -> GetReturn:
            ...
    
    
    def class_factory(s: GetReturn) -> type[ExampleClass[GetReturn]]:
        class SpecificClass(ExampleClass):
            def __init__(self):
                pass
    
            def get(self) -> GetReturn:
                return s
    
        return SpecificClass
    
    ASubclass = class_factory("A") 
    
    def f(x: ExampleClass[Literal["A"]]):
        pass
    
    f(ASubclass())
    

    Personally I would probably use Solution 1 in most cases, but how specific you need the type hint depends on your use case. Hope this helps!