Fully-typed Python decorator for functions, methods, staticmethods, and classmethods.

decorator.py

"""
Showcase a fully-typed decorator that can be applied to functions, methods,
staticmethods, and classmethods.

This example has some limitations due to the limited expressiveness of the
Python type system:

1. When applying the decorator to a function whose first argument is an
   optional type object, the returned function's first argument is shown as
   required.
2. It's not possible to apply the decorator to a staticmethod whose first
   argument is a type object which is a superclass of the class it's defined
   in, or whose first argument can hold an instance of the class it's defined
   in. The returned function's type hints will be incorrect.

It's possible to work around both of these limitations by removing the first
overload from the decorator, however doing so will require all classmethods to
use the included classdecorator instead of the normal decorator.
"""

import types
import unittest
from typing import (
    Any,
    Callable,
    Concatenate,
    Generic,
    ParamSpec,
    TypeVar,
    overload,
    reveal_type,
)

P = ParamSpec("P")
BoundP = ParamSpec("BoundP")
R_co = TypeVar("R_co", covariant=True)
S = TypeVar("S")
T = TypeVar("T", bound=type)


class Decorator(Generic[P, R_co]):
    def __init__(self, f: Callable[P, R_co]):
        self.f = f

    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R_co:
        return self.f(*args, **kwargs)

    @overload
    def __get__(
        self: "Decorator[Concatenate[S, BoundP], R_co]",
        instance: S,
        owner: type | None,
    ) -> "Decorator[BoundP, R_co]": ...
    @overload
    def __get__(self, instance: Any, owner: type | None) -> "Decorator[P, R_co]": ...
    def __get__(self, instance: Any, owner: Any = None) -> Any:
        # Overload 1 is for accessing through an instance (binding).
        # Overload 2 is for accessing through the class itself (non-binding).
        # We special case support for staticmethod and classmethod here.
        if isinstance(self.f, staticmethod) or isinstance(self.f, classmethod):
            return self.f.__get__(instance, owner)
        if instance is None:
            return self
        return Decorator(types.MethodType(self.f, instance))


# This class is only used when type checking, it does not exist at runtime. It
# is responsible for functions which take type objects as the first parameter,
# and handles the self binding.
class MethodDecorator(Generic[T, P, R_co]):

    def __call__(self, instance: T, *args: P.args, **kwargs: P.kwargs) -> R_co: ...

    # Overload 1 matches member accesses where the class on the left is the
    # same as the method's first parameter (classmethods). Overload 2 matches
    # everything else (assume non-binding).
    @overload
    def __get__(self, instance: Any, owner: T) -> "Decorator[P, R_co]": ...
    @overload
    def __get__(
        self, instance: Any, owner: type | None
    ) -> "Decorator[Concatenate[T, P], R_co]": ...
    def __get__(self, instance: Any, owner: type | None) -> Any: ...


@overload
def decorator(f: Callable[Concatenate[T, P], R_co]) -> MethodDecorator[T, P, R_co]: ...
@overload
def decorator(  # pyright: ignore[reportOverlappingOverload]
    f: Callable[P, R_co]
) -> Decorator[P, R_co]: ...
def decorator(f: Any) -> Any:
    return Decorator(f)


def classdecorator(f: Callable[Concatenate[T, P], R_co]) -> MethodDecorator[T, P, R_co]:
    return Decorator(f)  # type: ignore


@decorator
def func() -> None:
    print("in func()")


@decorator
def func_param(val: int = 1) -> None:
    print(f"in func_param({val})")


@decorator
def func_typevar(val: type[float]) -> None:
    print("in func_typevar")


class Class:
    @decorator
    def method(self) -> None:
        assert isinstance(self, Class)
        print("in Class.method()")

    @decorator
    def method_param(self, val: int = 1) -> None:
        assert isinstance(self, Class)
        print(f"in Class.method_param({val})")

    @decorator
    @staticmethod
    def static_method() -> None:
        print("in Class.static_method()")

    @decorator
    @staticmethod
    def static_method_param(val: int = 1) -> None:
        print(f"in Class.static_method_param({val})")

    @decorator
    @staticmethod
    def static_method_typevar(val: type[float]) -> None:
        print(f"in Class.static_method_typevar({val})")

    @decorator
    @staticmethod
    def static_method_typevar_same_type(val: "type[Class]") -> None:
        print(f"in Class.static_method_typevar_same_type({val})")

    @decorator
    @staticmethod
    def static_method_instance(val: object) -> None:
        print(f"in Class.static_method_instance({val})")

    @decorator
    @classmethod
    def class_method(cls) -> None:
        assert cls is Class
        print("in Class.class_method()")

    @decorator
    @classmethod
    def class_method_param(cls, val: int = 1) -> None:
        assert cls is Class
        print(f"in Class.class_method_param({val})")


class TestCases(unittest.TestCase):
    def test_runtime(self) -> None:
        reveal_type(func)
        func()

        reveal_type(func_param)
        func_param()
        func_param(1)

        reveal_type(func_typevar)
        func_typevar(int)

        reveal_type(Class.method)
        reveal_type(Class().method)
        Class.method(Class())
        Class().method()

        reveal_type(Class.method_param)
        reveal_type(Class().method_param)
        Class.method_param(Class(), 1)
        Class().method_param(1)
        Class.method_param(Class())
        Class().method_param()

        reveal_type(Class.static_method)
        reveal_type(Class().static_method)
        Class.static_method()
        Class().static_method()

        reveal_type(Class.static_method_param)
        reveal_type(Class().static_method_param)
        Class.static_method_param(1)
        Class().static_method_param(1)
        Class.static_method_param()
        Class().static_method_param()

        reveal_type(Class.static_method_typevar)
        reveal_type(Class().static_method_typevar)
        Class.static_method_typevar(int)
        Class().static_method_typevar(int)

        reveal_type(Class.class_method)
        reveal_type(Class().class_method)
        Class.class_method()
        Class().class_method()

        reveal_type(Class.class_method_param)
        reveal_type(Class().class_method_param)
        Class.class_method_param(1)
        Class().class_method_param(1)
        Class.class_method_param()
        Class().class_method_param()

    def test_typing_failures(self) -> None:
        reveal_type(Class.static_method_instance)
        reveal_type(Class().static_method_instance)
        Class.static_method_instance(Class())  # type: ignore
        Class().static_method_instance(Class())  # type: ignore

        reveal_type(Class.static_method_typevar_same_type)
        reveal_type(Class().static_method_typevar_same_type)
        Class.static_method_typevar_same_type(Class)  # type: ignore
        Class().static_method_typevar_same_type(Class)  # type: ignore


if __name__ == "__main__":
    unittest.main()

I've updated the gist. Changes in revision 4:

• Default arguments are supported for all arguments except those which are the first positional argument and can hold a type object (e.g. def func(first: object = SOME_OBJECT) is still not supported, but def func(first: int = 1) is).
• Functions which can hold a type object in their first argument are now supported (e.g. def func(first: object) is now supported).
• Staticmethods which can hold a type object in their first argument are now supported, except those which can also hold the class they're defined in (e.g. def static_method(first: object) is not supported, but def static_method(first: type[SomeUnrelatedClass]) is).
• Staticmethods which can hold an object in their first argument are now supported, except those which can also hold an instance of the class they're defined in (e.g. def static_method(first: object) is not supported, but def static_method(first: SomeUnrelatedClass) is).