Search code examples
pythonmypypython-typing

How to properly type hint a class decorator?


Assume we have some function func that maps instances of class A to instances of class B, i.e. it has the signature Callable[[A], B].

I want to write a class decorator autofunc for subclasses of A that automatically applies func to instances as they are created. For example, think of automatic jit-compilation based on a global environment variable. This can be done via

from functools import wraps

def autofunc(basecls):
    @wraps(basecls, updated=())

    class WrappedClass(basecls):
        def __new__(cls, *args, **kwargs):
            instance = basecls(*args, **kwargs)
            return func(instance)
    return WrappedClass

Then the following two are roughly equivalent:

class C(A):
...

instance = func(C())

@autofunc
class C(A):
...

instance = C()

In my naivety, I tried

def autofunc(basecls: type[A]) -> type[B]:

    @wraps(basecls, updated=())
    class WrappedClass(basecls):

        def __new__(cls, *args, **kwargs):
            instance = basecls(*args, **kwargs)
            return func(instance)

    return WrappedClass

which mypy really dislikes, raising erros:

  • error: Variable "basecls" is not valid as a type [valid-type]
  • error: Invalid base class "basecls" [misc]
  • also, there is the issue that WrapperClass.__new__ returns an instance of B and not of WrapperClass.

Is there any way to properly type hint such a class decorator currently, or is mypy not capable to work this yet?


Example code:
from functools import wraps


class A:
    pass


class B:
    pass


def func(cl: A) -> B:
    print(f"Replacing {cl=}")
    return B()


def autofunc(basecls: type[A]) -> type[B]:

    @wraps(basecls, updated=())
    class WrappedClass(basecls):

        def __new__(cls, *args, **kwargs):
            instance = basecls()
            return func(instance)

    return WrappedClass

Solution

  • This is simply a bug in mypy: https://github.com/python/mypy/issues/5865