10
10
from collections .abc import Iterator
11
11
from functools import partial
12
12
13
- from astroid import context , extract_node , inference_tip
13
+ from astroid import context , extract_node , inference_tip , nodes
14
14
from astroid .builder import _extract_single_node
15
15
from astroid .const import PY38_PLUS , PY39_PLUS
16
16
from astroid .exceptions import (
35
35
from astroid .util import Uninferable
36
36
37
37
TYPING_NAMEDTUPLE_BASENAMES = {"NamedTuple" , "typing.NamedTuple" }
38
- TYPING_TYPEVARS = {"TypeVar" , "NewType" }
39
- TYPING_TYPEVARS_QUALIFIED = {"typing.TypeVar" , "typing.NewType" }
40
38
TYPING_TYPE_TEMPLATE = """
41
39
class Meta(type):
42
40
def __getitem__(self, item):
@@ -49,6 +47,13 @@ def __args__(self):
49
47
class {0}(metaclass=Meta):
50
48
pass
51
49
"""
50
+ # PEP484 suggests NewType is equivalent to this for typing purposes
51
+ # https://www.python.org/dev/peps/pep-0484/#newtype-helper-function
52
+ TYPING_NEWTYPE_TEMPLATE = """
53
+ class {derived}({base}):
54
+ def __init__(self, val: {base}) -> None:
55
+ ...
56
+ """
52
57
TYPING_MEMBERS = set (getattr (typing , "__all__" , []))
53
58
54
59
TYPING_ALIAS = frozenset (
@@ -103,23 +108,34 @@ def __class_getitem__(cls, item):
103
108
"""
104
109
105
110
106
- def looks_like_typing_typevar_or_newtype (node ):
111
+ def looks_like_typing_typevar (node : nodes .Call ) -> bool :
112
+ func = node .func
113
+ if isinstance (func , Attribute ):
114
+ return func .attrname == "TypeVar"
115
+ if isinstance (func , Name ):
116
+ return func .name == "TypeVar"
117
+ return False
118
+
119
+
120
+ def looks_like_typing_newtype (node : nodes .Call ) -> bool :
107
121
func = node .func
108
122
if isinstance (func , Attribute ):
109
- return func .attrname in TYPING_TYPEVARS
123
+ return func .attrname == "NewType"
110
124
if isinstance (func , Name ):
111
- return func .name in TYPING_TYPEVARS
125
+ return func .name == "NewType"
112
126
return False
113
127
114
128
115
- def infer_typing_typevar_or_newtype (node , context_itton = None ):
116
- """Infer a typing.TypeVar(...) or typing.NewType(...) call"""
129
+ def infer_typing_typevar (
130
+ node : nodes .Call , ctx : context .InferenceContext | None = None
131
+ ) -> Iterator [nodes .ClassDef ]:
132
+ """Infer a typing.TypeVar(...) call"""
117
133
try :
118
- func = next (node .func .infer (context = context_itton ))
134
+ func = next (node .func .infer (context = ctx ))
119
135
except (InferenceError , StopIteration ) as exc :
120
136
raise UseInferenceDefault from exc
121
137
122
- if func .qname () not in TYPING_TYPEVARS_QUALIFIED :
138
+ if func .qname () != "typing.TypeVar" :
123
139
raise UseInferenceDefault
124
140
if not node .args :
125
141
raise UseInferenceDefault
@@ -129,7 +145,55 @@ def infer_typing_typevar_or_newtype(node, context_itton=None):
129
145
130
146
typename = node .args [0 ].as_string ().strip ("'" )
131
147
node = extract_node (TYPING_TYPE_TEMPLATE .format (typename ))
132
- return node .infer (context = context_itton )
148
+ return node .infer (context = ctx )
149
+
150
+
151
+ def infer_typing_newtype (
152
+ node : nodes .Call , ctx : context .InferenceContext | None = None
153
+ ) -> Iterator [nodes .ClassDef ]:
154
+ """Infer a typing.NewType(...) call"""
155
+ try :
156
+ func = next (node .func .infer (context = ctx ))
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
+ # Cannot infer from a dynamic class name (f-string)
166
+ if isinstance (node .args [0 ], JoinedStr ):
167
+ raise UseInferenceDefault
168
+
169
+ derived , base = node .args
170
+ derived_name = derived .as_string ().strip ("'" )
171
+ base_name = base .as_string ().strip ("'" )
172
+
173
+ new_node : ClassDef = extract_node (
174
+ TYPING_NEWTYPE_TEMPLATE .format (derived = derived_name , base = base_name )
175
+ )
176
+ new_node .parent = node .parent
177
+
178
+ # Base type arg is a normal reference, so no need to do special lookups
179
+ if not isinstance (base , nodes .Const ):
180
+ new_node .postinit (
181
+ bases = [base ], body = new_node .body , decorators = new_node .decorators
182
+ )
183
+
184
+ # If the base type is given as a string (e.g. for a forward reference),
185
+ # make a naive attempt to find the corresponding node.
186
+ # Note that this will not work with imported types.
187
+ if isinstance (base , nodes .Const ) and isinstance (base .value , str ):
188
+ _ , resolved_base = node .frame ().lookup (base_name )
189
+ if resolved_base :
190
+ new_node .postinit (
191
+ bases = [resolved_base [0 ]],
192
+ body = new_node .body ,
193
+ decorators = new_node .decorators ,
194
+ )
195
+
196
+ return new_node .infer (context = ctx )
133
197
134
198
135
199
def _looks_like_typing_subscript (node ):
@@ -403,8 +467,13 @@ def infer_typing_cast(
403
467
404
468
AstroidManager ().register_transform (
405
469
Call ,
406
- inference_tip (infer_typing_typevar_or_newtype ),
407
- looks_like_typing_typevar_or_newtype ,
470
+ inference_tip (infer_typing_typevar ),
471
+ looks_like_typing_typevar ,
472
+ )
473
+ AstroidManager ().register_transform (
474
+ Call ,
475
+ inference_tip (infer_typing_newtype ),
476
+ looks_like_typing_newtype ,
408
477
)
409
478
AstroidManager ().register_transform (
410
479
Subscript , inference_tip (infer_typing_attr ), _looks_like_typing_subscript
0 commit comments