How to type a decorator which also injects an attribute?

[BrandonNav]

Hello,

I have a decorator which injects an attribute onto a function:

from collections.abc import Callable
from functools import wraps
from typing import ParamSpec, TypeVar


_R = TypeVar("_R")
_P = ParamSpec("_P")


def mark_exempt(func: Callable[_P, _R]) -> Callable[_P, _R]:
    @wraps(func)
    def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R:
        return func(*args, **kwargs)

    setattr(wrapped, "exempt", True)
    return wrapped

The signature of mark_exempt is stating that the functions takes in a callable and returns a callable of the same signature which is technically correct. But the returned function object is a different type as it has gained the new exempt attribute.

I would like the type system to track the type change so that my IDE autocompletes the attribute as well as the type system being able to error if the decorator is missing.

My current attempt to solve this is as follows:

from collections.abc import Callable
from functools import wraps
from typing import Literal, ParamSpec, Protocol, TypeVar, reveal_type


_R = TypeVar("_R", covariant=True)
_P = ParamSpec("_P")


class CallableExemptForCSRF(Protocol[_P, _R, _P, _R]):
    exempt: Literal[True]

    def __call__(self, *args: _P.args, **kwds: _P.kwargs) -> _R:
        ...


def mark_exempt(func: Callable[_P, _R]) -> CallableExemptForCSRF[_P, _R]:
    @wraps(func)
    def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R:
        return func(*args, **kwargs)

    setattr(wrapped, "exempt", True)
    return wrapped  # This line errors in type checkers

Pyright errors due to protocol incompatibility (exempt is not present)
Mypy errors due to Incompatible return value type (got "Callable[_P, _R]", expected "CallableExemptForCSRF[_P, _R]")

If I ignore the pyright/mypy errors for return wrapped it does "work"

@mark_exempt
def a():
    pass

reveal_type(a.exempt)
# Revealed type is "Literal[True]"

I don't think a # type:ignore is the way to go here. I could cast to the protocol but I tend to avoid casting as I prefer not to bypass the typechecker. An alternative is marking the protocol as runtime_checkable and adding assert isinstance(wrapped_view, CallableExemptForCSRF) before the return but I am unsure if that is correct (or just equivalent to the cast).

How would I type this correctly without errors?

[erictraut]

You're modifying the value in a way that's not type safe, so you'd need to use a cast call.

Code sample in pyright playground

from functools import wraps
from typing import Callable, Literal, ParamSpec, Protocol, TypeVar, cast

_R = TypeVar("_R", covariant=True)
_P = ParamSpec("_P")

class CallableExemptForCSRF(Protocol[_P, _R]):
    exempt: Literal[True]

    def __call__(self, *args: _P.args, **kwds: _P.kwargs) -> _R:
        ...

def mark_exempt(func: Callable[_P, _R]) -> CallableExemptForCSRF[_P, _R]:
    @wraps(func)
    def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R:
        return func(*args, **kwargs)

    setattr(wrapped, "exempt", True)
    return cast(CallableExemptForCSRF[_P, _R], wrapped)

[BrandonNav]

Thanks for your reply @erictraut

So, if I'm understanding you correctly, any dynamic attribute modification is type unsafe. I imagine this is because descriptors can intercept the setattr and so you can't guarantee anything.

With this in mind and thinking about this more, since I can't modify a function in this way, I think a type-safe why to do this is using class with a __call__

eg:
Code sample in pyright playground

from collections.abc import Callable
from functools import update_wrapper
from typing import ClassVar, Generic, Literal, ParamSpec, TypeVar


_R = TypeVar("_R", covariant=True)
_P = ParamSpec("_P")


class CallableExemptForCSRF(Generic[_P, _R]):
    exempt: ClassVar[Literal[True]] = True

    def __init__(self, func: Callable[_P, _R]):
        update_wrapper(self, func)
        self._func = func

    def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R:
        return self._func(*args, **kwargs)


def mark_exempt(func: Callable[_P, _R]) -> CallableExemptForCSRF[_P, _R]:
    return CallableExemptForCSRF[_P, _R](func)