Trying to add type annotations to this decorator that adds an attribute to a function

I’ve got this code extracted and simplified from my schemafunc library:

from functools import wraps

def function_to_schema(func, **schema_kwargs):
    return {}

class FunctionMetadata:
    def __init__(self, func, **schema_kwargs):
        self.func = func
        self.schema_kwargs = schema_kwargs

    def schema(self):
        return function_to_schema(self.func, **self.schema_kwargs)

def add_schemafunc(_func=None, **schema_kwargs):
    def decorator(func):
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)

        if not hasattr(wrapper, "function_metadata"):
            wrapper.schemafunc = FunctionMetadata(func, **schema_kwargs)

        return wrapper

    if _func is not None:
        return decorator(_func)
        return decorator

I’d really like to add comprehensive type hints so that users can get typing on their_decorated_function.schemafunc. However, I’m running into problems trying to get this done.

I’ve gotten to the point where I kinda-sorta get it working. And by “working” I mean, mypy doesn’t complain about anything in function_to_schema, FunctionMetadata and add_schemafunc by using ParamSpec, TypeVar and a Protocol called HasSchemaFuncAttribute, but it falls apart when I actually try to use the decorator.

Here’s the kinda-sorta working version on the mypy playground.

You need to add overloads to your decorator/decorator combo function: mypy Playground