15
15
import typing
16
16
from functools import partial
17
17
18
- from astroid import context , extract_node , inference_tip
18
+ from astroid import context , extract_node , inference_tip , nodes
19
19
from astroid .const import PY37_PLUS , PY38_PLUS , PY39_PLUS
20
20
from astroid .exceptions import (
21
21
AttributeInferenceError ,
38
38
from astroid .util import Uninferable
39
39
40
40
TYPING_NAMEDTUPLE_BASENAMES = {"NamedTuple" , "typing.NamedTuple" }
41
- TYPING_TYPEVARS = {"TypeVar" , "NewType" }
42
- TYPING_TYPEVARS_QUALIFIED = {"typing.TypeVar" , "typing.NewType" }
43
41
TYPING_TYPE_TEMPLATE = """
44
42
class Meta(type):
45
43
def __getitem__(self, item):
@@ -52,6 +50,13 @@ def __args__(self):
52
50
class {0}(metaclass=Meta):
53
51
pass
54
52
"""
53
+ # PEP484 suggests NewType is equivalent to this for typing purposes
54
+ # https://www.python.org/dev/peps/pep-0484/#newtype-helper-function
55
+ TYPING_NEWTYPE_TEMPLATE = """
56
+ class {derived}({base}):
57
+ def __init__(self, val: {base}) -> None:
58
+ ...
59
+ """
55
60
TYPING_MEMBERS = set (getattr (typing , "__all__" , []))
56
61
57
62
TYPING_ALIAS = frozenset (
@@ -106,23 +111,34 @@ def __class_getitem__(cls, item):
106
111
"""
107
112
108
113
109
- def looks_like_typing_typevar_or_newtype (node ):
114
+ def looks_like_typing_typevar (node : nodes .Call ) -> bool :
115
+ func = node .func
116
+ if isinstance (func , Attribute ):
117
+ return func .attrname == "TypeVar"
118
+ if isinstance (func , Name ):
119
+ return func .name == "TypeVar"
120
+ return False
121
+
122
+
123
+ def looks_like_typing_newtype (node : nodes .Call ) -> bool :
110
124
func = node .func
111
125
if isinstance (func , Attribute ):
112
- return func .attrname in TYPING_TYPEVARS
126
+ return func .attrname == "NewType"
113
127
if isinstance (func , Name ):
114
- return func .name in TYPING_TYPEVARS
128
+ return func .name == "NewType"
115
129
return False
116
130
117
131
118
- def infer_typing_typevar_or_newtype (node , context_itton = None ):
119
- """Infer a typing.TypeVar(...) or typing.NewType(...) call"""
132
+ def infer_typing_typevar (
133
+ node : nodes .Call , context_itton : typing .Optional [context .InferenceContext ] = None
134
+ ) -> typing .Iterator [nodes .ClassDef ]:
135
+ """Infer a typing.TypeVar(...) call"""
120
136
try :
121
137
func = next (node .func .infer (context = context_itton ))
122
138
except (InferenceError , StopIteration ) as exc :
123
139
raise UseInferenceDefault from exc
124
140
125
- if func .qname () not in TYPING_TYPEVARS_QUALIFIED :
141
+ if func .qname () != "typing.TypeVar" :
126
142
raise UseInferenceDefault
127
143
if not node .args :
128
144
raise UseInferenceDefault
@@ -132,6 +148,44 @@ def infer_typing_typevar_or_newtype(node, context_itton=None):
132
148
return node .infer (context = context_itton )
133
149
134
150
151
+ def infer_typing_newtype (
152
+ node : nodes .Call , context_itton : typing .Optional [context .InferenceContext ] = None
153
+ ) -> typing .Iterator [nodes .ClassDef ]:
154
+ """Infer a typing.NewType(...) call"""
155
+ try :
156
+ func = next (node .func .infer (context = context_itton ))
157
+ except (InferenceError , StopIteration ) as exc :
158
+ raise UseInferenceDefault from exc
159
+
160
+ if func .qname () != "typing.NewType" :
161
+ raise UseInferenceDefault
162
+ if len (node .args ) != 2 :
163
+ raise UseInferenceDefault
164
+
165
+ derived , base = node .args
166
+ derived_name = derived .as_string ().strip ("'" )
167
+ base_name = base .as_string ().strip ("'" )
168
+
169
+ new_node : ClassDef = extract_node (
170
+ TYPING_NEWTYPE_TEMPLATE .format (derived = derived_name , base = base_name )
171
+ )
172
+ new_node .parent = node .parent
173
+
174
+ # Base type arg is a normal reference, so no need to do special lookups
175
+ if not isinstance (base , nodes .Const ):
176
+ new_node .bases = [base ]
177
+
178
+ # If the base type is given as a string (e.g. for a forward reference),
179
+ # make a naive attempt to find the corresponding node.
180
+ # Note that this will not work with imported types.
181
+ if isinstance (base , nodes .Const ) and isinstance (base .value , str ):
182
+ _ , resolved_base = node .frame ().lookup (base_name )
183
+ if resolved_base :
184
+ new_node .bases = [resolved_base [0 ]]
185
+
186
+ return new_node .infer (context = context_itton )
187
+
188
+
135
189
def _looks_like_typing_subscript (node ):
136
190
"""Try to figure out if a Subscript node *might* be a typing-related subscript"""
137
191
if isinstance (node , Name ):
@@ -409,8 +463,13 @@ def infer_typing_cast(
409
463
410
464
AstroidManager ().register_transform (
411
465
Call ,
412
- inference_tip (infer_typing_typevar_or_newtype ),
413
- looks_like_typing_typevar_or_newtype ,
466
+ inference_tip (infer_typing_typevar ),
467
+ looks_like_typing_typevar ,
468
+ )
469
+ AstroidManager ().register_transform (
470
+ Call ,
471
+ inference_tip (infer_typing_newtype ),
472
+ looks_like_typing_newtype ,
414
473
)
415
474
AstroidManager ().register_transform (
416
475
Subscript , inference_tip (infer_typing_attr ), _looks_like_typing_subscript
0 commit comments