-
Notifications
You must be signed in to change notification settings - Fork 6
/
type_hints.py
140 lines (114 loc) · 4.98 KB
/
type_hints.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
import inspect
import sys
import types
from collections.abc import Callable
from typing import Any, Union, cast, get_args, get_origin, get_type_hints
NoneType = cast(type, type(None)) # mypy otherwise treats type(None) as an object
def _check_issubclass(klass: Any, check_type: type) -> bool:
klass_origin, check_type_origin = get_origin(klass), get_origin(check_type)
klass_args, check_type_args = get_args(klass), get_args(check_type)
# eg: issubclass(tuple, tuple)
if klass_origin is None and check_type_origin is None:
return issubclass(klass, check_type)
# eg: issubclass(tuple[int], tuple)
if klass_origin is not None and check_type_origin is None:
return issubclass(klass_origin, check_type)
# eg: issubclass(tuple, tuple[int])
if klass_origin is None and check_type_origin is not None:
return issubclass(klass, check_type_origin) and not check_type_args
# eg: issubclass(tuple[int], tuple[int])
if klass_origin is not None and check_type_origin is not None:
# NOTE: Considering all container types covariant for simplicity (mypy may be more strict).
#
# The builtin mutable containers (list, dict, etc) are invariant (klass_args ==
# check_type_args), but the interfaces (Mapping, Sequence, etc) and immutable containers are
# covariant.
if check_type_args:
if not (
len(klass_args) == len(check_type_args)
and all(
lenient_issubclass(klass_arg, check_type_arg)
for (klass_arg, check_type_arg) in zip(klass_args, check_type_args)
)
):
return False
return lenient_issubclass(klass_origin, check_type_origin)
# Shouldn't happen, but need to explicitly say "x is not None" to narrow mypy types.
raise NotImplementedError("The origin conditions don't cover all cases!")
def get_class_type_vars(klass: type) -> tuple[type, ...]:
base = klass.__orig_bases__[0] # type: ignore
base_origin = get_origin(base)
if base_origin is None:
raise TypeError(f"{klass.__name__} must subclass a subscripted Generic")
assert lenient_issubclass(klass, base_origin)
return get_args(base)
def lenient_issubclass(klass: Any, class_or_tuple: Union[type, tuple[type, ...]]) -> bool:
if not isinstance(klass, type):
return False
if isinstance(class_or_tuple, tuple):
return any(lenient_issubclass(klass, subtype) for subtype in class_or_tuple)
check_type = class_or_tuple
if is_union_hint(check_type):
return any(lenient_issubclass(klass, subtype) for subtype in get_args(check_type))
return _check_issubclass(klass, check_type)
def _tidy_return(return_annotation: Any, *, force_tuple_return: bool) -> Any:
if not force_tuple_return:
return return_annotation
if lenient_issubclass(get_origin(return_annotation), tuple):
return get_args(return_annotation)
return (return_annotation,)
def tidy_signature(
fn: Callable[..., Any],
sig: inspect.Signature,
*,
force_tuple_return: bool = False,
remove_owner: bool = False,
) -> inspect.Signature:
type_hints = get_type_hints(fn)
sig = sig.replace(return_annotation=type_hints.get("return", sig.return_annotation))
return sig.replace(
parameters=[
p.replace(annotation=type_hints.get(p.name, p.annotation))
for p in sig.parameters.values()
if (p.name not in ("cls", "self") if remove_owner else True)
],
return_annotation=(
sig.empty
if sig.return_annotation is sig.empty
else _tidy_return(sig.return_annotation, force_tuple_return=force_tuple_return)
),
)
def signature(
fn: Callable[..., Any],
*,
follow_wrapped: bool = True,
force_tuple_return: bool = False,
remove_owner: bool = True,
) -> inspect.Signature:
"""Convenience wrapper around `inspect.signature`.
The returned Signature will have `cls`/`self` parameters removed if `remove_owner` is `True` and
`tuple[...]` converted to `tuple(...)` in the `return_annotation`.
"""
return tidy_signature(
fn=fn,
sig=inspect.signature(fn, follow_wrapped=follow_wrapped),
force_tuple_return=force_tuple_return,
remove_owner=remove_owner,
)
#############################################
# Helpers for typing across python versions #
#############################################
#
# Focusing on 3.9+ (for now)
if sys.version_info < (3, 10):
def is_union(type_: Any) -> bool:
return type_ is Union
else: # pragma: no cover
def is_union(type_: Any) -> bool:
# `Union[int, str]` or `int | str`
return type_ is Union or type_ is types.UnionType # noqa: E721
def is_optional_hint(type_: Any) -> bool:
# Optional[x] is represented as Union[x, NoneType]
return is_union(get_origin(type_)) and NoneType in get_args(type_)
def is_union_hint(type_: Any) -> bool:
return get_origin(type_) is Union