brain_typing.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438
  1. # Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
  2. # For details: https://github.com/PyCQA/astroid/blob/main/LICENSE
  3. # Copyright (c) 2017-2018 Claudiu Popa <pcmanticore@gmail.com>
  4. # Copyright (c) 2017 Łukasz Rogalski <rogalski.91@gmail.com>
  5. # Copyright (c) 2017 David Euresti <github@euresti.com>
  6. # Copyright (c) 2018 Bryce Guinta <bryce.paul.guinta@gmail.com>
  7. # Copyright (c) 2021 Pierre Sassoulas <pierre.sassoulas@gmail.com>
  8. # Copyright (c) 2021 Daniël van Noord <13665637+DanielNoord@users.noreply.github.com>
  9. # Copyright (c) 2021 Redoubts <Redoubts@users.noreply.github.com>
  10. # Copyright (c) 2021 Marc Mueller <30130371+cdce8p@users.noreply.github.com>
  11. # Copyright (c) 2021 Tim Martin <tim@asymptotic.co.uk>
  12. # Copyright (c) 2021 hippo91 <guillaume.peillex@gmail.com>
  13. """Astroid hooks for typing.py support."""
  14. import typing
  15. from functools import partial
  16. from astroid import context, extract_node, inference_tip
  17. from astroid.const import PY37_PLUS, PY38_PLUS, PY39_PLUS
  18. from astroid.exceptions import (
  19. AttributeInferenceError,
  20. InferenceError,
  21. UseInferenceDefault,
  22. )
  23. from astroid.manager import AstroidManager
  24. from astroid.nodes.node_classes import (
  25. Assign,
  26. AssignName,
  27. Attribute,
  28. Call,
  29. Const,
  30. Name,
  31. NodeNG,
  32. Subscript,
  33. Tuple,
  34. )
  35. from astroid.nodes.scoped_nodes import ClassDef, FunctionDef
  36. from astroid.util import Uninferable
  37. TYPING_NAMEDTUPLE_BASENAMES = {"NamedTuple", "typing.NamedTuple"}
  38. TYPING_TYPEVARS = {"TypeVar", "NewType"}
  39. TYPING_TYPEVARS_QUALIFIED = {"typing.TypeVar", "typing.NewType"}
  40. TYPING_TYPE_TEMPLATE = """
  41. class Meta(type):
  42. def __getitem__(self, item):
  43. return self
  44. @property
  45. def __args__(self):
  46. return ()
  47. class {0}(metaclass=Meta):
  48. pass
  49. """
  50. TYPING_MEMBERS = set(getattr(typing, "__all__", []))
  51. TYPING_ALIAS = frozenset(
  52. (
  53. "typing.Hashable",
  54. "typing.Awaitable",
  55. "typing.Coroutine",
  56. "typing.AsyncIterable",
  57. "typing.AsyncIterator",
  58. "typing.Iterable",
  59. "typing.Iterator",
  60. "typing.Reversible",
  61. "typing.Sized",
  62. "typing.Container",
  63. "typing.Collection",
  64. "typing.Callable",
  65. "typing.AbstractSet",
  66. "typing.MutableSet",
  67. "typing.Mapping",
  68. "typing.MutableMapping",
  69. "typing.Sequence",
  70. "typing.MutableSequence",
  71. "typing.ByteString",
  72. "typing.Tuple",
  73. "typing.List",
  74. "typing.Deque",
  75. "typing.Set",
  76. "typing.FrozenSet",
  77. "typing.MappingView",
  78. "typing.KeysView",
  79. "typing.ItemsView",
  80. "typing.ValuesView",
  81. "typing.ContextManager",
  82. "typing.AsyncContextManager",
  83. "typing.Dict",
  84. "typing.DefaultDict",
  85. "typing.OrderedDict",
  86. "typing.Counter",
  87. "typing.ChainMap",
  88. "typing.Generator",
  89. "typing.AsyncGenerator",
  90. "typing.Type",
  91. "typing.Pattern",
  92. "typing.Match",
  93. )
  94. )
  95. CLASS_GETITEM_TEMPLATE = """
  96. @classmethod
  97. def __class_getitem__(cls, item):
  98. return cls
  99. """
  100. def looks_like_typing_typevar_or_newtype(node):
  101. func = node.func
  102. if isinstance(func, Attribute):
  103. return func.attrname in TYPING_TYPEVARS
  104. if isinstance(func, Name):
  105. return func.name in TYPING_TYPEVARS
  106. return False
  107. def infer_typing_typevar_or_newtype(node, context_itton=None):
  108. """Infer a typing.TypeVar(...) or typing.NewType(...) call"""
  109. try:
  110. func = next(node.func.infer(context=context_itton))
  111. except (InferenceError, StopIteration) as exc:
  112. raise UseInferenceDefault from exc
  113. if func.qname() not in TYPING_TYPEVARS_QUALIFIED:
  114. raise UseInferenceDefault
  115. if not node.args:
  116. raise UseInferenceDefault
  117. typename = node.args[0].as_string().strip("'")
  118. node = extract_node(TYPING_TYPE_TEMPLATE.format(typename))
  119. return node.infer(context=context_itton)
  120. def _looks_like_typing_subscript(node):
  121. """Try to figure out if a Subscript node *might* be a typing-related subscript"""
  122. if isinstance(node, Name):
  123. return node.name in TYPING_MEMBERS
  124. if isinstance(node, Attribute):
  125. return node.attrname in TYPING_MEMBERS
  126. if isinstance(node, Subscript):
  127. return _looks_like_typing_subscript(node.value)
  128. return False
  129. def infer_typing_attr(
  130. node: Subscript, ctx: typing.Optional[context.InferenceContext] = None
  131. ) -> typing.Iterator[ClassDef]:
  132. """Infer a typing.X[...] subscript"""
  133. try:
  134. value = next(node.value.infer())
  135. except (InferenceError, StopIteration) as exc:
  136. raise UseInferenceDefault from exc
  137. if (
  138. not value.qname().startswith("typing.")
  139. or PY37_PLUS
  140. and value.qname() in TYPING_ALIAS
  141. ):
  142. # If typing subscript belongs to an alias
  143. # (PY37+) handle it separately.
  144. raise UseInferenceDefault
  145. if (
  146. PY37_PLUS
  147. and isinstance(value, ClassDef)
  148. and value.qname()
  149. in {"typing.Generic", "typing.Annotated", "typing_extensions.Annotated"}
  150. ):
  151. # With PY37+ typing.Generic and typing.Annotated (PY39) are subscriptable
  152. # through __class_getitem__. Since astroid can't easily
  153. # infer the native methods, replace them for an easy inference tip
  154. func_to_add = extract_node(CLASS_GETITEM_TEMPLATE)
  155. value.locals["__class_getitem__"] = [func_to_add]
  156. if (
  157. isinstance(node.parent, ClassDef)
  158. and node in node.parent.bases
  159. and getattr(node.parent, "__cache", None)
  160. ):
  161. # node.parent.slots is evaluated and cached before the inference tip
  162. # is first applied. Remove the last result to allow a recalculation of slots
  163. cache = node.parent.__cache # type: ignore[attr-defined] # Unrecognized getattr
  164. if cache.get(node.parent.slots) is not None:
  165. del cache[node.parent.slots]
  166. return iter([value])
  167. node = extract_node(TYPING_TYPE_TEMPLATE.format(value.qname().split(".")[-1]))
  168. return node.infer(context=ctx)
  169. def _looks_like_typedDict( # pylint: disable=invalid-name
  170. node: typing.Union[FunctionDef, ClassDef],
  171. ) -> bool:
  172. """Check if node is TypedDict FunctionDef."""
  173. return node.qname() in {"typing.TypedDict", "typing_extensions.TypedDict"}
  174. def infer_old_typedDict( # pylint: disable=invalid-name
  175. node: ClassDef, ctx: typing.Optional[context.InferenceContext] = None
  176. ) -> typing.Iterator[ClassDef]:
  177. func_to_add = extract_node("dict")
  178. node.locals["__call__"] = [func_to_add]
  179. return iter([node])
  180. def infer_typedDict( # pylint: disable=invalid-name
  181. node: FunctionDef, ctx: typing.Optional[context.InferenceContext] = None
  182. ) -> typing.Iterator[ClassDef]:
  183. """Replace TypedDict FunctionDef with ClassDef."""
  184. class_def = ClassDef(
  185. name="TypedDict",
  186. lineno=node.lineno,
  187. col_offset=node.col_offset,
  188. parent=node.parent,
  189. )
  190. class_def.postinit(bases=[extract_node("dict")], body=[], decorators=None)
  191. func_to_add = extract_node("dict")
  192. class_def.locals["__call__"] = [func_to_add]
  193. return iter([class_def])
  194. def _looks_like_typing_alias(node: Call) -> bool:
  195. """
  196. Returns True if the node corresponds to a call to _alias function.
  197. For example :
  198. MutableSet = _alias(collections.abc.MutableSet, T)
  199. :param node: call node
  200. """
  201. return (
  202. isinstance(node.func, Name)
  203. and node.func.name == "_alias"
  204. and (
  205. # _alias function works also for builtins object such as list and dict
  206. isinstance(node.args[0], (Attribute, Name))
  207. )
  208. )
  209. def _forbid_class_getitem_access(node: ClassDef) -> None:
  210. """
  211. Disable the access to __class_getitem__ method for the node in parameters
  212. """
  213. def full_raiser(origin_func, attr, *args, **kwargs):
  214. """
  215. Raises an AttributeInferenceError in case of access to __class_getitem__ method.
  216. Otherwise just call origin_func.
  217. """
  218. if attr == "__class_getitem__":
  219. raise AttributeInferenceError("__class_getitem__ access is not allowed")
  220. return origin_func(attr, *args, **kwargs)
  221. try:
  222. node.getattr("__class_getitem__")
  223. # If we are here, then we are sure to modify object that do have __class_getitem__ method (which origin is one the
  224. # protocol defined in collections module) whereas the typing module consider it should not
  225. # We do not want __class_getitem__ to be found in the classdef
  226. partial_raiser = partial(full_raiser, node.getattr)
  227. node.getattr = partial_raiser
  228. except AttributeInferenceError:
  229. pass
  230. def infer_typing_alias(
  231. node: Call, ctx: typing.Optional[context.InferenceContext] = None
  232. ) -> typing.Iterator[ClassDef]:
  233. """
  234. Infers the call to _alias function
  235. Insert ClassDef, with same name as aliased class,
  236. in mro to simulate _GenericAlias.
  237. :param node: call node
  238. :param context: inference context
  239. """
  240. if (
  241. not isinstance(node.parent, Assign)
  242. or not len(node.parent.targets) == 1
  243. or not isinstance(node.parent.targets[0], AssignName)
  244. ):
  245. raise UseInferenceDefault
  246. try:
  247. res = next(node.args[0].infer(context=ctx))
  248. except StopIteration as e:
  249. raise InferenceError(node=node.args[0], context=context) from e
  250. assign_name = node.parent.targets[0]
  251. class_def = ClassDef(
  252. name=assign_name.name,
  253. lineno=assign_name.lineno,
  254. col_offset=assign_name.col_offset,
  255. parent=node.parent,
  256. )
  257. if res != Uninferable and isinstance(res, ClassDef):
  258. # Only add `res` as base if it's a `ClassDef`
  259. # This isn't the case for `typing.Pattern` and `typing.Match`
  260. class_def.postinit(bases=[res], body=[], decorators=None)
  261. maybe_type_var = node.args[1]
  262. if (
  263. not PY39_PLUS
  264. and not (isinstance(maybe_type_var, Tuple) and not maybe_type_var.elts)
  265. or PY39_PLUS
  266. and isinstance(maybe_type_var, Const)
  267. and maybe_type_var.value > 0
  268. ):
  269. # If typing alias is subscriptable, add `__class_getitem__` to ClassDef
  270. func_to_add = extract_node(CLASS_GETITEM_TEMPLATE)
  271. class_def.locals["__class_getitem__"] = [func_to_add]
  272. else:
  273. # If not, make sure that `__class_getitem__` access is forbidden.
  274. # This is an issue in cases where the aliased class implements it,
  275. # but the typing alias isn't subscriptable. E.g., `typing.ByteString` for PY39+
  276. _forbid_class_getitem_access(class_def)
  277. return iter([class_def])
  278. def _looks_like_special_alias(node: Call) -> bool:
  279. """Return True if call is for Tuple or Callable alias.
  280. In PY37 and PY38 the call is to '_VariadicGenericAlias' with 'tuple' as
  281. first argument. In PY39+ it is replaced by a call to '_TupleType'.
  282. PY37: Tuple = _VariadicGenericAlias(tuple, (), inst=False, special=True)
  283. PY39: Tuple = _TupleType(tuple, -1, inst=False, name='Tuple')
  284. PY37: Callable = _VariadicGenericAlias(collections.abc.Callable, (), special=True)
  285. PY39: Callable = _CallableType(collections.abc.Callable, 2)
  286. """
  287. return isinstance(node.func, Name) and (
  288. not PY39_PLUS
  289. and node.func.name == "_VariadicGenericAlias"
  290. and (
  291. isinstance(node.args[0], Name)
  292. and node.args[0].name == "tuple"
  293. or isinstance(node.args[0], Attribute)
  294. and node.args[0].as_string() == "collections.abc.Callable"
  295. )
  296. or PY39_PLUS
  297. and (
  298. node.func.name == "_TupleType"
  299. and isinstance(node.args[0], Name)
  300. and node.args[0].name == "tuple"
  301. or node.func.name == "_CallableType"
  302. and isinstance(node.args[0], Attribute)
  303. and node.args[0].as_string() == "collections.abc.Callable"
  304. )
  305. )
  306. def infer_special_alias(
  307. node: Call, ctx: typing.Optional[context.InferenceContext] = None
  308. ) -> typing.Iterator[ClassDef]:
  309. """Infer call to tuple alias as new subscriptable class typing.Tuple."""
  310. if not (
  311. isinstance(node.parent, Assign)
  312. and len(node.parent.targets) == 1
  313. and isinstance(node.parent.targets[0], AssignName)
  314. ):
  315. raise UseInferenceDefault
  316. try:
  317. res = next(node.args[0].infer(context=ctx))
  318. except StopIteration as e:
  319. raise InferenceError(node=node.args[0], context=context) from e
  320. assign_name = node.parent.targets[0]
  321. class_def = ClassDef(
  322. name=assign_name.name,
  323. parent=node.parent,
  324. )
  325. class_def.postinit(bases=[res], body=[], decorators=None)
  326. func_to_add = extract_node(CLASS_GETITEM_TEMPLATE)
  327. class_def.locals["__class_getitem__"] = [func_to_add]
  328. return iter([class_def])
  329. def _looks_like_typing_cast(node: Call) -> bool:
  330. return isinstance(node, Call) and (
  331. isinstance(node.func, Name)
  332. and node.func.name == "cast"
  333. or isinstance(node.func, Attribute)
  334. and node.func.attrname == "cast"
  335. )
  336. def infer_typing_cast(
  337. node: Call, ctx: typing.Optional[context.InferenceContext] = None
  338. ) -> typing.Iterator[NodeNG]:
  339. """Infer call to cast() returning same type as casted-from var"""
  340. if not isinstance(node.func, (Name, Attribute)):
  341. raise UseInferenceDefault
  342. try:
  343. func = next(node.func.infer(context=ctx))
  344. except (InferenceError, StopIteration) as exc:
  345. raise UseInferenceDefault from exc
  346. if (
  347. not isinstance(func, FunctionDef)
  348. or func.qname() != "typing.cast"
  349. or len(node.args) != 2
  350. ):
  351. raise UseInferenceDefault
  352. return node.args[1].infer(context=ctx)
  353. AstroidManager().register_transform(
  354. Call,
  355. inference_tip(infer_typing_typevar_or_newtype),
  356. looks_like_typing_typevar_or_newtype,
  357. )
  358. AstroidManager().register_transform(
  359. Subscript, inference_tip(infer_typing_attr), _looks_like_typing_subscript
  360. )
  361. AstroidManager().register_transform(
  362. Call, inference_tip(infer_typing_cast), _looks_like_typing_cast
  363. )
  364. if PY39_PLUS:
  365. AstroidManager().register_transform(
  366. FunctionDef, inference_tip(infer_typedDict), _looks_like_typedDict
  367. )
  368. elif PY38_PLUS:
  369. AstroidManager().register_transform(
  370. ClassDef, inference_tip(infer_old_typedDict), _looks_like_typedDict
  371. )
  372. if PY37_PLUS:
  373. AstroidManager().register_transform(
  374. Call, inference_tip(infer_typing_alias), _looks_like_typing_alias
  375. )
  376. AstroidManager().register_transform(
  377. Call, inference_tip(infer_special_alias), _looks_like_special_alias
  378. )