1818import typing
1919from functools import partial
2020
21- from astroid import context , extract_node , inference_tip
21+ from astroid import context , extract_node , inference_tip , nodes
2222from astroid .builder import _extract_single_node
2323from astroid .const import PY37_PLUS , PY38_PLUS , PY39_PLUS
2424from astroid .exceptions import (
4343from astroid .util import Uninferable
4444
4545TYPING_NAMEDTUPLE_BASENAMES = {"NamedTuple" , "typing.NamedTuple" }
46- TYPING_TYPEVARS = {"TypeVar" , "NewType" }
47- TYPING_TYPEVARS_QUALIFIED = {"typing.TypeVar" , "typing.NewType" }
4846TYPING_TYPE_TEMPLATE = """
4947class Meta(type):
5048 def __getitem__(self, item):
@@ -57,6 +55,13 @@ def __args__(self):
5755class {0}(metaclass=Meta):
5856 pass
5957"""
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+ """
6065TYPING_MEMBERS = set (getattr (typing , "__all__" , []))
6166
6267TYPING_ALIAS = frozenset (
@@ -111,23 +116,34 @@ def __class_getitem__(cls, item):
111116"""
112117
113118
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 :
115129 func = node .func
116130 if isinstance (func , Attribute ):
117- return func .attrname in TYPING_TYPEVARS
131+ return func .attrname == "NewType"
118132 if isinstance (func , Name ):
119- return func .name in TYPING_TYPEVARS
133+ return func .name == "NewType"
120134 return False
121135
122136
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"""
125141 try :
126142 func = next (node .func .infer (context = context_itton ))
127143 except (InferenceError , StopIteration ) as exc :
128144 raise UseInferenceDefault from exc
129145
130- if func .qname () not in TYPING_TYPEVARS_QUALIFIED :
146+ if func .qname () != "typing.TypeVar" :
131147 raise UseInferenceDefault
132148 if not node .args :
133149 raise UseInferenceDefault
@@ -140,6 +156,54 @@ def infer_typing_typevar_or_newtype(node, context_itton=None):
140156 return node .infer (context = context_itton )
141157
142158
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 .postinit (
189+ bases = [base ], body = new_node .body , decorators = new_node .decorators
190+ )
191+
192+ # If the base type is given as a string (e.g. for a forward reference),
193+ # make a naive attempt to find the corresponding node.
194+ # Note that this will not work with imported types.
195+ if isinstance (base , nodes .Const ) and isinstance (base .value , str ):
196+ _ , resolved_base = node .frame ().lookup (base_name )
197+ if resolved_base :
198+ new_node .postinit (
199+ bases = [resolved_base [0 ]],
200+ body = new_node .body ,
201+ decorators = new_node .decorators ,
202+ )
203+
204+ return new_node .infer (context = context_itton )
205+
206+
143207def _looks_like_typing_subscript (node ):
144208 """Try to figure out if a Subscript node *might* be a typing-related subscript"""
145209 if isinstance (node , Name ):
@@ -417,8 +481,13 @@ def infer_typing_cast(
417481
418482AstroidManager ().register_transform (
419483 Call ,
420- inference_tip (infer_typing_typevar_or_newtype ),
421- looks_like_typing_typevar_or_newtype ,
484+ inference_tip (infer_typing_typevar ),
485+ looks_like_typing_typevar ,
486+ )
487+ AstroidManager ().register_transform (
488+ Call ,
489+ inference_tip (infer_typing_newtype ),
490+ looks_like_typing_newtype ,
422491)
423492AstroidManager ().register_transform (
424493 Subscript , inference_tip (infer_typing_attr ), _looks_like_typing_subscript
0 commit comments