Search code examples
pythontype-hintingtyping

Restricted generic type hint in Python


I want to create a factory method like here.

class A:
    ...

class B(A):
    def f(self):
        ...

class C:
    ...

def factory(cls):
    return cls()

But I want to add some type hints with two requirements:

  • Only subclasses of A are allowed as arguments of factory.
  • When B is passed, it is correcly detected that factory(B) is instance of B, ie. factory(B).f() is allowed while factory(A).f() is not.

Attempt 1

Use Type from typing module.

from typing import Type

class A:
    ...

class B(A):
    def f(self):
        ...

class C:
    ...

def factory(cls: Type[A]):
    return cls()

factory(A)  # Ok
factory(B)  # Ok
factory(C)  # Fail
factory(A).f()  # Fail
factory(B).f()  # Fail -- Wrong!

This one correctly detects that C is not suppose to be passed as argument of factory. However, factory(B).f() is not allowed by the type-checker.

Attempt 2

Use TypeVar.

from typing import TypeVar, Type
T = TypeVar('T')

class A:
    ...

class B(A):
    def f(self):
        ...

class C:
    ...

def factory(cls: Type[T]) -> T:
    return cls()

factory(A)  # Ok
factory(B)  # Ok
factory(C)  # Ok -- Wrong!
factory(A).f()  # Fail
factory(B).f()  # Ok

It is well inferred that factory(B).f() is fine, while factory(A).f() is not. However, the generics is without restriction, ie. factory(C) is Ok as well.

Attempt 3

Add constraint to T.

from typing import TypeVar, Type

class A:
    ...

class B(A):
    def f(self):
        ...

class D(A):
    ...

class C:
    ...

T = TypeVar('T', A)
def factory(cls: Type[T]) -> T:
    return cls()


factory(A)  # Ok
factory(B)  # Ok
factory(C)  # Fail
factory(A).f()  # Fail
factory(B).f()  # Ok

This looks promising and at least PyCharm properly handles all cases. But for practical use is this solution actually the worst - A single constraint is not allowed error arises, though, it is not exactly clear why. In PEP 484, there is only a brief line saying 'There should be at least two constraints, if any; specifying a single constraint is disallowed.'

Is there any nice solution to this? I have only something like adding a 'dummy' class _A, a blank subclass of A, and putting it as another constraint to have something like this.

_A = NewType('_A', A)
T = TypeVar('T', A, _A)
def factory(cls: Type[T]) -> T:
    return cls()

But I don't really consider it as a nice solution.

Thank you in advance!


Solution

  • Using a TypeVar with a bound seems to work:

    from typing import Type, TypeVar
    
    class A:
        ...
    
    class B(A):
        def f(self):
            ...
    
    class C:
        ...
    
    AnyA = TypeVar("AnyA", bound=A)
    
    def factory(cls: Type[AnyA]) -> AnyA:
        return cls()
    
    
    factory(A).f()  # error: "A" has no attribute "f"
    factory(B).f()  # ok!
    factory(C)      # error: Value of type variable "AnyA" of "factory" cannot be "C"