Python 3 type hinting for decorator which changes argument types

89 Views Asked by At

I have a decorator which coerces function arguments into their type hinted types:

import inspect
from functools import wraps
from typing import Any, Callable, TypeVar

R = TypeVar("R")


def coerce_arguments(func: Callable[..., R]) -> Callable[..., R]:
    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> R:
        signature = inspect.signature(func)
        bound_args = signature.bind(*args, **kwargs)
        bound_args.apply_defaults()

        new_args = []
        for name, param in signature.parameters.items():
            if name in bound_args.arguments:
                value = bound_args.arguments[name]
                if param.annotation != inspect.Parameter.empty:
                    try:
                        coerced_value = param.annotation(value)
                    except Exception:
                        # attempt to run the function un-coerced
                        coerced_value = value
                    bound_args.arguments[name] = coerced_value
            new_args.append(bound_args.arguments[name])

        return func(*new_args, **kwargs)

    return wrapper


@coerce_arguments
def test(x: int) -> int:
    return x + 1


test("1")
# >>> 2

This works fine, but now my type checker will allow me to call test with any number of arguments ((...) -> int) - not very helpful. Instead, the function should still be called with the same number of arguments, but with Any type for all these arguments. I.e. (x: Any) -> int for test.

I've tried using typing.ParamSpec with no success :(

import inspect
from functools import wraps
from typing import Any, Callable, TypeVar, ParamSpec

R = TypeVar("R")
P = ParamSpec("P")
O = ParamSpec("O")


def coerce_arguments(func: Callable[P, R]) -> Callable[O, R]:
    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> R:  # Not sure what types for args and kwargs
        ...
        new_args = []  # type hint new_args?
        ...
        return func(*new_args, **kwargs)

    return wrapper


@coerce_arguments
def test(x: int) -> int:
    return x + 1

# (**O@coerce_arguments) -> int
1

There are 1 best solutions below

1
hungry_in_learning On

To address this issue, you can use typing.get_type_hints to obtain the type hints for the function parameters and then apply coercion only if the parameter has a concrete type hint. With the document

import inspect
from functools import wraps
from typing import Any, Callable, TypeVar, get_type_hints

R = TypeVar("R")
P = TypeVar("P", bound=Callable[..., Any])
O = TypeVar("O", bound=Callable[..., Any])


def coerce_arguments(func: P) -> O:
    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> R:
        # Get the type hints for the function parameters
        type_hints = get_type_hints(func)

        new_args = []
        for name, value in zip(func.__code__.co_varnames, args):
            param_hint = type_hints.get(name, Any)
            if param_hint != Any:
                try:
                    coerced_value = param_hint(value)
                except (TypeError, ValueError):
                    # Failed to coerce, use the original value
                    coerced_value = value
                new_args.append(coerced_value)
            else:
                new_args.append(value)

        return func(*new_args, **kwargs)

    return wrapper


@coerce_arguments
def test(x: int) -> int:
    return x + 1


test("1")  # Output: 2