Subclass in type hinting

127.2k views Asked by At

I want to allow type hinting using Python 3 to accept sub classes of a certain class. E.g.:

class A:
    pass

class B(A):
    pass

class C(A):
    pass

def process_any_subclass_type_of_A(cls: A):
    if cls == B:
        # do something
    elif cls == C:
        # do something else

Now when typing the following code:

process_any_subclass_type_of_A(B)

I get an PyCharm IDE hint 'Expected type A, got Type[B] instead.'

How can I change type hinting here to accept any subtypes of A?

According to PEP 484 ("Expressions whose type is a subtype of a specific argument type are also accepted for that argument."), I understand that my solution (cls: A) should work?

3

There are 3 answers

4
Ashwini Chaudhary On BEST ANSWER

When you specify cls: A, you're saying that cls expects an instance of type A. The type hint to specify cls as a class object for the type A (or its subtypes) uses typing.Type.

from typing import Type
def process_any_subclass_type_of_A(cls: Type[A]):
    pass

From The type of class objects :

Sometimes you want to talk about class objects that inherit from a given class. This can be spelled as Type[C] where C is a class. In other words, when C is the name of a class, using C to annotate an argument declares that the argument is an instance of C (or of a subclass of C), but using Type[C] as an argument annotation declares that the argument is a class object deriving from C (or C itself).

4
FRiMN On

If we look at the Type description from the typing module, then we see these docs:

A special construct usable to annotate class objects.

For example, suppose we have the following classes::

 class User: ...  # Abstract base for User classes
 class BasicUser(User): ...
 class ProUser(User): ...
 class TeamUser(User): ...

And a function that takes a class argument that's a subclass of User and returns an instance of the corresponding class::

 U = TypeVar('U', bound=User)
 def new_user(user_class: Type[U]) -> U:
     user = user_class()
     # (Here we could write the user object to a database)
     return user

 joe = new_user(BasicUser)

At this point the type checker knows that joe has type BasicUser.

Based on this, I can imagine a synthetic example that reproduces the problem with type hinting errors in PyCharm.

from typing import Type, Tuple

class BaseClass: ...

class SubClass(BaseClass): ...
class SubSubClass(SubClass): ...

def process(model_instance: BaseClass, model_class: Type[BaseClass]) -> Tuple[BaseClass, BaseClass]:
    """ Accepts all of the above classes """
    return model_instance, model_class()


class ProcessorA:
    @staticmethod
    def proc() -> Tuple[SubClass, SubClass]:
        """ PyCharm will show an error 
        `Expected type 'tuple[SubClass, SubClass]', got 'tuple[BaseClass, BaseClass]' instead` """
        return process(SubClass(), SubClass)


class ProcessorB:
    @staticmethod
    def proc() -> Tuple[SubSubClass, SubSubClass]:
        """ PyCharm will show an error 
        `Expected type 'tuple[SubSubClass, SubSubClass]', got 'tuple[BaseClass, BaseClass]' instead` """
        return process(SubSubClass(), SubSubClass)

But we see in docs for Type that the situation can be corrected by using TypeVar with the bound argument. Then use it in places where BaseClass is declared as a type.

from typing import TypeVar, Type, Tuple

class BaseClass: ...

B = TypeVar('B', bound=BaseClass)

class SubClass(BaseClass): ...
class SubSubClass(SubClass): ...

def process(model_instance: B, model_class: Type[B]) -> Tuple[B, B]:
    """ Accepts all of the above classes """
    return model_instance, model_class()


class ProcessorA:
    @staticmethod
    def proc() -> Tuple[SubClass, SubClass]:
        return process(SubClass(), SubClass)


class ProcessorB:
    @staticmethod
    def proc() -> Tuple[SubSubClass, SubSubClass]:
        return process(SubSubClass(), SubSubClass)

Hope this will be helpful.

1
BouygudSt On

Type[A] accepts also the class itself, which is not always needed.

If you want your function to accept only subclasses, you should go with NewType, like

class A:
    pass

B = NewType('B', A)

def foo(cls: Type[B]):
   ...