CODEWITHSUNDEEP
pythonpython-typingpython-decorators
MikeL
MikeL

Reputation: 2469

How to correct the decorated function signature and type hints?

Consider the following decorator that extends any binary operator to multiple arguments:

from typing import Callable, TypeVar
from functools import reduce, wraps


T = TypeVar('T')


def extend(binop: Callable[[T, T], T]):
    """ Extend a binary operator to multiple arguments """

    @wraps(binop)
    def extended(*args: T) -> T:
        if not args:
            raise TypeError("At least one argument must be given")

        return reduce(binop, args)

    return extended

which may then be used as follows:

@extend
def fadd(x: float, y: float) -> float:
    """ Add float numbers """

    return x + y


@extend
def imul(x: int, y: int) -> int:
    """ Multiply integers """

    return x*y

so as to create the imul and fadd functions that multiply and add their input arguments respectively.

The functions imul and fadd will have the correct docstrings (because of @wraps decorator) but their signatures and type annotations are incorrect. For example:

>>> help(fadd)

Gives

fadd(x: float, y: float) -> float                                               
    Add float numbers

Also

>>> fadd.__annotations__                                              
{'x': <class 'float'>, 'y': <class 'float'>, 'return': <class 'float'>}

which is incorrect.

What is the correct way to implement a decorator to result in the right function signature?

I somehow thought that if I remove @wraps line the type hints and signature will be correct. But even that is not the case. Without @wraps

>>> help(fadd)

gives

extended(*args: ~T) -> ~T

(i.e., the generic type T is not replaced by float).

Upvotes: 5

Views: 906

Answers (1)

MikeL
MikeL

Reputation: 2469

I found a slightly hacky solution using the inspect

def extend(binop: Callable[[T, T], T]):
    """ Extend a binary operator to multiple arguments """

    @wraps(binop)
    def extended(*args: T) -> T:
        if not args:
            raise TypeError("At least one argument must be given")

        return reduce(binop, args)

    sig = inspect.signature(extended)

    sig = sig.replace(
        parameters=[inspect.Parameter('args', inspect.Parameter.VAR_POSITIONAL,
                                      annotation=sig.return_annotation)]
    )

    extended.__signature__ = sig

    return extended

It is not as elegant as I was expecting, especially because it heavily depends on the fact that binop return values must have the same type as its arguments (otherwise the annotations will be wrong).

Yet, I would be happy to know of better solutions.

Upvotes: 1

Related Questions