18
18
import typing
19
19
from functools import partial
20
20
21
- from astroid import context , extract_node , inference_tip
21
+ from astroid import context , extract_node , inference_tip , nodes
22
22
from astroid .builder import _extract_single_node
23
23
from astroid .const import PY37_PLUS , PY38_PLUS , PY39_PLUS
24
24
from astroid .exceptions import (
43
43
from astroid .util import Uninferable
44
44
45
45
TYPING_NAMEDTUPLE_BASENAMES = {"NamedTuple" , "typing.NamedTuple" }
46
- TYPING_TYPEVARS = {"TypeVar" , "NewType" }
47
- TYPING_TYPEVARS_QUALIFIED = {"typing.TypeVar" , "typing.NewType" }
48
46
TYPING_TYPE_TEMPLATE = """
49
47
class Meta(type):
50
48
def __getitem__(self, item):
@@ -57,6 +55,13 @@ def __args__(self):
57
55
class {0}(metaclass=Meta):
58
56
pass
59
57
"""
58
+ # PEP484 suggests NewType is equivalent to this for typing purposes
59
+ # https://www.python.org/dev/peps/pep-0484/#newtype-helper-function
60
+ TYPING_NEWTYPE_TEMPLATE = """
61
+ class {derived}({base}):
62
+ def __init__(self, val: {base}) -> None:
63
+ ...
64
+ """
60
65
TYPING_MEMBERS = set (getattr (typing , "__all__" , []))
61
66
62
67
TYPING_ALIAS = frozenset (
@@ -111,23 +116,34 @@ def __class_getitem__(cls, item):
111
116
"""
112
117
113
118
114
- def looks_like_typing_typevar_or_newtype (node ):
119
+ def looks_like_typing_typevar (node : nodes .Call ) -> bool :
120
+ func = node .func
121
+ if isinstance (func , Attribute ):
122
+ return func .attrname == "TypeVar"
123
+ if isinstance (func , Name ):
124
+ return func .name == "TypeVar"
125
+ return False
126
+
127
+
128
+ def looks_like_typing_newtype (node : nodes .Call ) -> bool :
115
129
func = node .func
116
130
if isinstance (func , Attribute ):
117
- return func .attrname in TYPING_TYPEVARS
131
+ return func .attrname == "NewType"
118
132
if isinstance (func , Name ):
119
- return func .name in TYPING_TYPEVARS
133
+ return func .name == "NewType"
120
134
return False
121
135
122
136
123
- def infer_typing_typevar_or_newtype (node , context_itton = None ):
124
- """Infer a typing.TypeVar(...) or typing.NewType(...) call"""
137
+ def infer_typing_typevar (
138
+ node : nodes .Call , context_itton : typing .Optional [context .InferenceContext ] = None
139
+ ) -> typing .Iterator [nodes .ClassDef ]:
140
+ """Infer a typing.TypeVar(...) call"""
125
141
try :
126
142
func = next (node .func .infer (context = context_itton ))
127
143
except (InferenceError , StopIteration ) as exc :
128
144
raise UseInferenceDefault from exc
129
145
130
- if func .qname () not in TYPING_TYPEVARS_QUALIFIED :
146
+ if func .qname () != "typing.TypeVar" :
131
147
raise UseInferenceDefault
132
148
if not node .args :
133
149
raise UseInferenceDefault
@@ -140,6 +156,48 @@ def infer_typing_typevar_or_newtype(node, context_itton=None):
140
156
return node .infer (context = context_itton )
141
157
142
158
159
+ def infer_typing_newtype (
160
+ node : nodes .Call , context_itton : typing .Optional [context .InferenceContext ] = None
161
+ ) -> typing .Iterator [nodes .ClassDef ]:
162
+ """Infer a typing.NewType(...) call"""
163
+ try :
164
+ func = next (node .func .infer (context = context_itton ))
165
+ except (InferenceError , StopIteration ) as exc :
166
+ raise UseInferenceDefault from exc
167
+
168
+ if func .qname () != "typing.NewType" :
169
+ raise UseInferenceDefault
170
+ if len (node .args ) != 2 :
171
+ raise UseInferenceDefault
172
+
173
+ # Cannot infer from a dynamic class name (f-string)
174
+ if isinstance (node .args [0 ], JoinedStr ):
175
+ raise UseInferenceDefault
176
+
177
+ derived , base = node .args
178
+ derived_name = derived .as_string ().strip ("'" )
179
+ base_name = base .as_string ().strip ("'" )
180
+
181
+ new_node : ClassDef = extract_node (
182
+ TYPING_NEWTYPE_TEMPLATE .format (derived = derived_name , base = base_name )
183
+ )
184
+ new_node .parent = node .parent
185
+
186
+ # Base type arg is a normal reference, so no need to do special lookups
187
+ if not isinstance (base , nodes .Const ):
188
+ new_node .bases = [base ]
189
+
190
+ # If the base type is given as a string (e.g. for a forward reference),
191
+ # make a naive attempt to find the corresponding node.
192
+ # Note that this will not work with imported types.
193
+ if isinstance (base , nodes .Const ) and isinstance (base .value , str ):
194
+ _ , resolved_base = node .frame ().lookup (base_name )
195
+ if resolved_base :
196
+ new_node .bases = [resolved_base [0 ]]
197
+
198
+ return new_node .infer (context = context_itton )
199
+
200
+
143
201
def _looks_like_typing_subscript (node ):
144
202
"""Try to figure out if a Subscript node *might* be a typing-related subscript"""
145
203
if isinstance (node , Name ):
@@ -417,8 +475,13 @@ def infer_typing_cast(
417
475
418
476
AstroidManager ().register_transform (
419
477
Call ,
420
- inference_tip (infer_typing_typevar_or_newtype ),
421
- looks_like_typing_typevar_or_newtype ,
478
+ inference_tip (infer_typing_typevar ),
479
+ looks_like_typing_typevar ,
480
+ )
481
+ AstroidManager ().register_transform (
482
+ Call ,
483
+ inference_tip (infer_typing_newtype ),
484
+ looks_like_typing_newtype ,
422
485
)
423
486
AstroidManager ().register_transform (
424
487
Subscript , inference_tip (infer_typing_attr ), _looks_like_typing_subscript
0 commit comments