1
|
|
|
import inspect |
2
|
|
|
import re |
3
|
|
|
from functools import wraps |
4
|
|
|
from typing import Dict, Optional, Callable, List |
5
|
|
|
|
6
|
|
|
_DEFAULT_PARAM_NAME = 'hint' |
7
|
|
|
|
8
|
|
|
|
9
|
|
View Code Duplication |
class _Hintable: |
|
|
|
|
10
|
|
|
_hints_per_frame = {} |
11
|
|
|
|
12
|
|
|
def __init__( |
13
|
|
|
self, |
14
|
|
|
decorated: Callable, |
15
|
|
|
param: str, |
16
|
|
|
stack_index: int) -> None: |
17
|
|
|
self._decorated = decorated |
18
|
|
|
self._param = param |
19
|
|
|
self._stack_index = stack_index |
20
|
|
|
|
21
|
|
|
def __call__(self, *args, **kwargs): |
22
|
|
|
stack = inspect.stack() |
23
|
|
|
|
24
|
|
|
previous_frame = stack[self._stack_index] |
25
|
|
|
frame_id = id(previous_frame.frame) |
26
|
|
|
|
27
|
|
|
if not self._hints_per_frame.get(frame_id): |
28
|
|
|
code_context = previous_frame.code_context[0].strip() |
29
|
|
|
hint_strs = self._extract_hints(code_context) |
30
|
|
|
globals_ = previous_frame.frame.f_globals |
31
|
|
|
# Store the type hint if any, otherwise the string, otherwise None. |
32
|
|
|
hints = [self._to_cls(hint_str, globals_) or hint_str or None |
33
|
|
|
for hint_str in hint_strs] |
34
|
|
|
self._hints_per_frame[frame_id] = hints |
35
|
|
|
|
36
|
|
|
hint = (self._hints_per_frame.get(frame_id) or [None]).pop() |
37
|
|
|
|
38
|
|
|
kwargs_ = {**kwargs, self._param: kwargs.get(self._param, hint)} |
39
|
|
|
return self._decorated(*args, **kwargs_) |
40
|
|
|
|
41
|
|
|
def _extract_hints(self, code_context: str) -> List[str]: |
42
|
|
|
result = [] |
43
|
|
|
regex = ( |
44
|
|
|
r'.+?(:(.+?))?=\s*' # e.g. 'x: int = ', $2 holds hint |
45
|
|
|
r'.*?{}\s*\(.*?\)\s*' # e.g. 'func(...) ' |
46
|
|
|
r'(#\s*type\s*:\s*(\w+))?\s*' # e.g. '# type: int ', $4 holds hint |
47
|
|
|
).format(self._decorated.__name__) |
48
|
|
|
|
49
|
|
|
# Find all matches and store them (reversed) in the resulting list. |
50
|
|
|
for _, group2, _, group4 in re.findall(regex, code_context): |
51
|
|
|
raw_hint = (group2 or group4).strip() |
52
|
|
|
if self._is_between(raw_hint, '\'') or self._is_between(raw_hint, '"'): |
53
|
|
|
# Remove any quotes that surround the hint. |
54
|
|
|
raw_hint = raw_hint[1:-1].strip() |
55
|
|
|
result.insert(0, raw_hint) |
56
|
|
|
|
57
|
|
|
return result |
58
|
|
|
|
59
|
|
|
def _is_between(self, subject: str, character: str) -> bool: |
60
|
|
|
return subject.startswith(character) and subject.endswith(character) |
61
|
|
|
|
62
|
|
|
def _to_cls(self, hint: str, f_globals: Dict[str, type]) -> Optional[type]: |
63
|
|
|
return __builtins__.get(hint, f_globals.get(hint)) |
|
|
|
|
64
|
|
|
|
65
|
|
|
|
66
|
|
View Code Duplication |
def _get_wrapper(decorated, param: str, stack_index: int): |
|
|
|
|
67
|
|
|
@wraps(decorated) |
68
|
|
|
def _wrapper(*args, **kwargs): |
69
|
|
|
return _Hintable(decorated, param, stack_index)(*args, **kwargs) |
70
|
|
|
|
71
|
|
|
if isinstance(decorated, type): |
72
|
|
|
raise TypeError('Only functions and methods should be decorated with ' |
73
|
|
|
'\'hintable\', not classes.') |
74
|
|
|
|
75
|
|
|
if param not in inspect.signature(decorated).parameters: |
76
|
|
|
raise TypeError('The decorated \'{}\' must accept a parameter with ' |
77
|
|
|
'the name \'{}\'.' |
78
|
|
|
.format(decorated.__name__, param)) |
79
|
|
|
|
80
|
|
|
return _wrapper |
81
|
|
|
|
82
|
|
|
|
83
|
|
View Code Duplication |
def hintable(decorated=None, *, param: str = _DEFAULT_PARAM_NAME) -> Callable: |
|
|
|
|
84
|
|
|
""" |
85
|
|
|
Allow a function or method to receive the type hint of a receiving |
86
|
|
|
variable. |
87
|
|
|
|
88
|
|
|
Example: |
89
|
|
|
|
90
|
|
|
>>> @hintable |
91
|
|
|
... def cast(value, hint): |
92
|
|
|
... return hint(value) |
93
|
|
|
>>> x: int = cast('42') |
94
|
|
|
42 |
95
|
|
|
|
96
|
|
|
Use this decorator wisely. If a variable was hinted with a type (e.g. int |
97
|
|
|
in the above example), your function should return a value of that type |
98
|
|
|
(in the above example, that would be an int value). |
99
|
|
|
|
100
|
|
|
:param decorated: a function or method. |
101
|
|
|
:param param: the name of the parameter that receives the type hint. |
102
|
|
|
:return: the decorated function/method wrapped into a new function. |
103
|
|
|
""" |
104
|
|
|
if decorated is not None: |
105
|
|
|
wrapper = _get_wrapper(decorated, param, 2) |
106
|
|
|
else: |
107
|
|
|
wrapper = lambda decorated_: _get_wrapper(decorated_, param, 2) |
108
|
|
|
|
109
|
|
|
return wrapper |
110
|
|
|
|