|
# -*- coding: UTF-8 -*- |
|
""" |
|
solve_overloads |
|
=============== |
|
|
|
Author |
|
------ |
|
Yuchen Jin |
|
[email protected] |
|
|
|
Description |
|
----------- |
|
Turn the `@overload` decorator from typehint-purposed to run-time effective. |
|
|
|
The module provides a `solve_overloads` decorator which is specifically designed for |
|
decorating function with overloads. The decorated function will be able to dispatch |
|
the input arguments according to the overload versions. |
|
|
|
The decorated function will try to bind the input arguments with each overload until |
|
a binding is successful. The trial order of the overloads follows the same order of |
|
the definition. |
|
""" |
|
|
|
import inspect |
|
import functools |
|
|
|
from typing import Any, TypeVar, cast |
|
from typing_extensions import Protocol, get_overloads |
|
|
|
|
|
class ArgsProtocol(Protocol): |
|
def __call__(self, *args: Any, **kwargs: Any) -> Any: ... |
|
|
|
|
|
_ArgsProtocol = TypeVar("_ArgsProtocol", bound=ArgsProtocol) |
|
__all__ = ("solve_overloads",) |
|
|
|
|
|
def solve_overloads(func: _ArgsProtocol) -> _ArgsProtocol: |
|
"""Decorator for branching overloads. |
|
|
|
This method can turn the typehint-purposed `@overload` decorator to an effective |
|
argument validator during the run time. |
|
|
|
This decorator is used for decorating any function with overloads. The decorator |
|
will find a wrapped function that tries to bind the provided arguments to each |
|
signature of the overloads until a succesful binding is done. If the arguments |
|
cannot be bounded to any overload, raise a `TypeError`. |
|
|
|
Arguments |
|
--------- |
|
func: `(*args: Any, **kwargs: Any) -> Any` |
|
The function to be bounded. The prototype of this function needs to contain |
|
`*args` and `**kwargs`. The function also needs to have at least one overload. |
|
|
|
Returns |
|
------- |
|
#1: `(*args: Any, **kwargs: Any) -> Any` |
|
The wrapped function. It has totally the same signature and the same |
|
functionality of the original function. However. After the decoration, |
|
the input arguments in the function implementation will be like this: |
|
```python |
|
{ |
|
"ver": int, |
|
"arguments": {*args, **kwargs} |
|
} |
|
``` |
|
Note that if the overload contains var arguments like `*args` or `**kwargs`, |
|
the values will be compressed as sequences or dictionaries, respectively. |
|
The value `ver` is the version of the overload. |
|
For example, |
|
```python |
|
@overload |
|
def func(val: int): ... |
|
|
|
|
|
@overload |
|
def func(*vals: int, **kwvals: str): ... |
|
|
|
|
|
@solve_overloads |
|
def func(*args, **kwargs): |
|
print(args) |
|
print(kwargs) |
|
|
|
|
|
func(1, 2, 3, val="a", val2="b") |
|
``` |
|
The printed results will be |
|
```python |
|
() |
|
{ |
|
"ver": 1, |
|
"arguments": { |
|
"vals": (1, 2, 3), |
|
"kwvals": {"val": "a", "val2": "b"} |
|
} |
|
} |
|
``` |
|
""" |
|
|
|
@functools.wraps(func) |
|
def wrapped(*args, **kwargs): |
|
for idx_ver, func_ver in enumerate(get_overloads(func)): |
|
try: |
|
all_args = inspect.signature(func_ver).bind(*args, **kwargs) |
|
except TypeError: |
|
continue |
|
all_args.apply_defaults() |
|
return func(ver=idx_ver, arguments=all_args.arguments) |
|
raise TypeError( |
|
"The provided arguments do not match any overload of the function." |
|
) |
|
|
|
return cast(_ArgsProtocol, wrapped) |
