pythonpython-typing

Decorator changes function structure


I'm facing a problem with annotating decorators. A raw function can take some arguments and a 'deps' argument. The decorator itself substitutes this argument. Question: how to annotate this decorator and will 'wraps' work correctly? Code:

P = ParamSpec("P")
T = TypeVar("T")

class Injector:
    _registered_dependencies: dict[str, Any] = {}
    def __init__(self, *dependency: str):
        self.dependency = dependency


    def __call__(self, func: Callable[Concatenate[dict[str, Any], P], T]) -> Callable[P, T]:
        @functools.wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            if "deps" in kwargs:
                raise ValueError("Argument 'deps' is already provided and will be overwritten.")
            try:
                deps = {
                    i: self.__class__._registered_dependencies[i] for i in self.dependency
                }
            except KeyError as e:
                raise KeyError(f"Dependency {e.args[0]} is not registered")
            return func(*args, **kwargs, deps=deps)
        return wrapper


    @classmethod
    def register(cls, id: str, dependency: Any) -> None:
        if not isinstance(id, str):
            raise TypeError("Dependency ID must be a string")
        if id in cls._registered_dependencies:
            raise ValueError(f"Dependency '{id}' is already registered")
        cls._registered_dependencies[id] = dependency


@Injector('test')
def foo(baz: str, deps: dict[str, Any]) -> None:
    print(f'{baz=}, {deps=}')  # baz='bar', deps={'test': 123}


if __name__ == '__main__':
    Injector.register('test', 423)
    foo('bar')

Tried using Concatenate[dict[str, Any], P] in __call__ types, also swapping arguments in Concatenate in places, but mypy still swears. Mypy says:

error: Argument 1 has incompatible type "*P.args"; expected "dict[str, Any]"  [arg-type]
error: Argument "deps" has incompatible type "dict[str, Any]"; expected "P.kwargs"  [arg-type]
error: Argument 1 to "__call__" of "Injector" has incompatible type "Callable[[str, dict[str, Any]], None]"; expected "Callable[[dict[str, Any], dict[str, Any]], None]"  [arg-type]
note: This is likely because "foo" has named arguments: "baz". Consider marking them positional-only
error: Argument 1 to "foo" has incompatible type "str"; expected "dict[str, Any]"  [arg-type]

Solution

  • Be aware that Callable[Concatenate[dict[str, Any], P], T] means the first argument must be a dict, all others are collected within P.

    That means the deps argument must also be the first argument, i.e.

    @Injector('test')
    def foo(deps: dict[str, Any], baz: str) -> None:
        print(f'{baz=}, {deps=}')  # baz='bar', deps={'test': 123}
    

    However you will face runtime problems when the argument is not named deps, this will not be reported by a type-checker, so you have to change to return func(deps, *args, **kwargs)


    If you want to force that the first name is called deps you use a Protocol instead:

    class HasDepsFirst(Protocol[P, T]):
    
        @staticmethod
        def __call__(deps: dict[str, Any], *args: P.args, **kwargs: P.kwargs) -> T: ...
    
        # alternatively similar to Callable this allows any name for the argument.
        # def __call__(deps: dict[str, Any], /, *args: P.args, **kwargs: P.kwargs) -> T: ...
    
    class Injector:
    
        def __call__(self, func: HasDepsFirst[P, T]) -> Callable[P, T]:
    

    Code sample in pyright playground