How to type hint a class decorator?

What is the correct way to annotate a decorator which takes a class as input and returns the same class as output?

For example:

def classdecorator(cls):
    cls.new_attrib = True
    return cls

@classdecorator
class MyClass:
    pass

MyClass.new_attrib
# True

What is the type of cls in classdecorator, and what is its return type?

2 Likes

You can treat it like an identity function:

def classdecorator[T](cls: T) -> T:
    cls.new_attrib = True
    return cls

If you want to prevent the decorator from being applied to things other than classes, you can bound the type var to type:

def classdecorator[T: type](cls: T) -> T:
    cls.new_attrib = True
    return cls

Or using pre PEP 695 syntax:

from typing import TypeVar

_T = TypeVar("_T", bound=type)

def classdecorator(cls: _T) -> _T:
    cls.new_attrib = True
    return cls
2 Likes

I think the goal, though, would be to assert that the type is altered, which would involve protocols. (The incoming class doesn’t necessarily have a new_attrib attribute, the resulting class does.)

Also, mypy at least rejects the attempt to create an attribute that an arbitrary type T doesn’t already have.

1 Like

Right, mypy complains about ‘“_T” has no attribute “new_attrib”’. Protocol does the trick:

class HasNewAttrib(Protocol):
    new_attrib: bool

_T = TypeVar("_T", bound=HasNewAttrib)

def classdecorator(cls: type[_T]) -> type[_T]:
    cls.new_attrib = True

though perhaps the input type should be something else…

It would have to be, because a class that doesn’t already have a new_attrib attribute wouldn’t be a valid argument.

1 Like

Hmm, which immediately unsolves the problem because then mypy thinks ‘type[_T] has no attribute “new_attrib”’ again. Which is correct, but the whole point of the decorator is to add that attribute.

If you want type checkers to be able to recognize that the return type is both compatible with the passed class as well as with HasNewAttrib, what you’d need is an intersection type. This is unfortunately a well known missing feature in Python’s type system. I think the best you can do for now is a union.

You can find some more info on the status of intersection types in links contained in this comment:

For more inspiration, you can take a look at how @dataclass is type hinted in typeshed. Note that even though @dataclass adds an attribute (__dataclass_fields__), its type hints can’t express that:

1 Like