1# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
2# For details: https://github.com/pylint-dev/astroid/blob/main/LICENSE
3# Copyright (c) https://github.com/pylint-dev/astroid/blob/main/CONTRIBUTORS.txt
4
5"""Astroid hooks for the Python standard library."""
6
7from __future__ import annotations
8
9import functools
10import keyword
11from collections.abc import Iterator
12from textwrap import dedent
13from typing import Final
14
15import astroid
16from astroid import arguments, bases, inference_tip, nodes, util
17from astroid.builder import AstroidBuilder, _extract_single_node, extract_node
18from astroid.context import InferenceContext
19from astroid.exceptions import (
20 AstroidTypeError,
21 AstroidValueError,
22 InferenceError,
23 UseInferenceDefault,
24)
25from astroid.manager import AstroidManager
26
27ENUM_QNAME: Final[str] = "enum.Enum"
28TYPING_NAMEDTUPLE_QUALIFIED: Final = {
29 "typing.NamedTuple",
30 "typing_extensions.NamedTuple",
31}
32TYPING_NAMEDTUPLE_BASENAMES: Final = {
33 "NamedTuple",
34 "typing.NamedTuple",
35 "typing_extensions.NamedTuple",
36}
37
38
39def _infer_first(node, context):
40 if isinstance(node, util.UninferableBase):
41 raise UseInferenceDefault
42 try:
43 value = next(node.infer(context=context))
44 except StopIteration as exc:
45 raise InferenceError from exc
46 if isinstance(value, util.UninferableBase):
47 raise UseInferenceDefault()
48 return value
49
50
51def _find_func_form_arguments(node, context):
52 def _extract_namedtuple_arg_or_keyword( # pylint: disable=inconsistent-return-statements
53 position, key_name=None
54 ):
55 if len(args) > position:
56 return _infer_first(args[position], context)
57 if key_name and key_name in found_keywords:
58 return _infer_first(found_keywords[key_name], context)
59
60 args = node.args
61 keywords = node.keywords
62 found_keywords = (
63 {keyword.arg: keyword.value for keyword in keywords} if keywords else {}
64 )
65
66 name = _extract_namedtuple_arg_or_keyword(position=0, key_name="typename")
67 names = _extract_namedtuple_arg_or_keyword(position=1, key_name="field_names")
68 if name and names:
69 return name.value, names
70
71 raise UseInferenceDefault()
72
73
74def infer_func_form(
75 node: nodes.Call,
76 base_type: list[nodes.NodeNG],
77 context: InferenceContext | None = None,
78 enum: bool = False,
79) -> tuple[nodes.ClassDef, str, list[str]]:
80 """Specific inference function for namedtuple or Python 3 enum."""
81 # node is a Call node, class name as first argument and generated class
82 # attributes as second argument
83
84 # namedtuple or enums list of attributes can be a list of strings or a
85 # whitespace-separate string
86 try:
87 name, names = _find_func_form_arguments(node, context)
88 try:
89 attributes: list[str] = names.value.replace(",", " ").split()
90 except AttributeError as exc:
91 # Handle attributes of NamedTuples
92 if not enum:
93 attributes = []
94 fields = _get_namedtuple_fields(node)
95 if fields:
96 fields_node = extract_node(fields)
97 attributes = [
98 _infer_first(const, context).value for const in fields_node.elts
99 ]
100
101 # Handle attributes of Enums
102 else:
103 # Enums supports either iterator of (name, value) pairs
104 # or mappings.
105 if hasattr(names, "items") and isinstance(names.items, list):
106 attributes = [
107 _infer_first(const[0], context).value
108 for const in names.items
109 if isinstance(const[0], nodes.Const)
110 ]
111 elif hasattr(names, "elts"):
112 # Enums can support either ["a", "b", "c"]
113 # or [("a", 1), ("b", 2), ...], but they can't
114 # be mixed.
115 if all(isinstance(const, nodes.Tuple) for const in names.elts):
116 attributes = [
117 _infer_first(const.elts[0], context).value
118 for const in names.elts
119 if isinstance(const, nodes.Tuple)
120 ]
121 else:
122 attributes = [
123 _infer_first(const, context).value for const in names.elts
124 ]
125 else:
126 raise AttributeError from exc
127 if not attributes:
128 raise AttributeError from exc
129 except (AttributeError, InferenceError) as exc:
130 raise UseInferenceDefault from exc
131
132 if not enum:
133 # namedtuple maps sys.intern(str()) over over field_names
134 attributes = [str(attr) for attr in attributes]
135 # XXX this should succeed *unless* __str__/__repr__ is incorrect or throws
136 # in which case we should not have inferred these values and raised earlier
137 attributes = [attr for attr in attributes if " " not in attr]
138
139 # If we can't infer the name of the class, don't crash, up to this point
140 # we know it is a namedtuple anyway.
141 name = name or "Uninferable"
142 # we want to return a Class node instance with proper attributes set
143 class_node = nodes.ClassDef(
144 name,
145 lineno=node.lineno,
146 col_offset=node.col_offset,
147 end_lineno=node.end_lineno,
148 end_col_offset=node.end_col_offset,
149 parent=nodes.Unknown(),
150 )
151 # A typical ClassDef automatically adds its name to the parent scope,
152 # but doing so causes problems, so defer setting parent until after init
153 # see: https://github.com/pylint-dev/pylint/issues/5982
154 class_node.parent = node.parent
155 class_node.postinit(
156 # set base class=tuple
157 bases=base_type,
158 body=[],
159 decorators=None,
160 )
161 # XXX add __init__(*attributes) method
162 for attr in attributes:
163 fake_node = nodes.EmptyNode()
164 fake_node.parent = class_node
165 fake_node.attrname = attr
166 class_node.instance_attrs[attr] = [fake_node]
167 return class_node, name, attributes
168
169
170def _has_namedtuple_base(node):
171 """Predicate for class inference tip.
172
173 :type node: ClassDef
174 :rtype: bool
175 """
176 return set(node.basenames) & TYPING_NAMEDTUPLE_BASENAMES
177
178
179def _looks_like(node, name) -> bool:
180 func = node.func
181 if isinstance(func, nodes.Attribute):
182 return func.attrname == name
183 if isinstance(func, nodes.Name):
184 return func.name == name
185 return False
186
187
188_looks_like_namedtuple = functools.partial(_looks_like, name="namedtuple")
189_looks_like_enum = functools.partial(_looks_like, name="Enum")
190_looks_like_typing_namedtuple = functools.partial(_looks_like, name="NamedTuple")
191
192
193def infer_named_tuple(
194 node: nodes.Call, context: InferenceContext | None = None
195) -> Iterator[nodes.ClassDef]:
196 """Specific inference function for namedtuple Call node."""
197 tuple_base_name: list[nodes.NodeNG] = [
198 nodes.Name(
199 name="tuple",
200 parent=node.root(),
201 lineno=0,
202 col_offset=0,
203 end_lineno=None,
204 end_col_offset=None,
205 )
206 ]
207 class_node, name, attributes = infer_func_form(
208 node, tuple_base_name, context=context
209 )
210 call_site = arguments.CallSite.from_call(node, context=context)
211 node = extract_node("import collections; collections.namedtuple")
212 try:
213 func = next(node.infer())
214 except StopIteration as e:
215 raise InferenceError(node=node) from e
216 try:
217 rename = next(
218 call_site.infer_argument(func, "rename", context or InferenceContext())
219 ).bool_value()
220 except (InferenceError, StopIteration):
221 rename = False
222
223 try:
224 attributes = _check_namedtuple_attributes(name, attributes, rename)
225 except AstroidTypeError as exc:
226 raise UseInferenceDefault("TypeError: " + str(exc)) from exc
227 except AstroidValueError as exc:
228 raise UseInferenceDefault("ValueError: " + str(exc)) from exc
229
230 replace_args = ", ".join(f"{arg}=None" for arg in attributes)
231 field_def = (
232 " {name} = property(lambda self: self[{index:d}], "
233 "doc='Alias for field number {index:d}')"
234 )
235 field_defs = "\n".join(
236 field_def.format(name=name, index=index)
237 for index, name in enumerate(attributes)
238 )
239 fake = AstroidBuilder(AstroidManager()).string_build(
240 f"""
241class {name}(tuple):
242 __slots__ = ()
243 _fields = {attributes!r}
244 def _asdict(self):
245 return self.__dict__
246 @classmethod
247 def _make(cls, iterable, new=tuple.__new__, len=len):
248 return new(cls, iterable)
249 def _replace(self, {replace_args}):
250 return self
251 def __getnewargs__(self):
252 return tuple(self)
253{field_defs}
254 """
255 )
256 class_node.locals["_asdict"] = fake.body[0].locals["_asdict"]
257 class_node.locals["_make"] = fake.body[0].locals["_make"]
258 class_node.locals["_replace"] = fake.body[0].locals["_replace"]
259 class_node.locals["_fields"] = fake.body[0].locals["_fields"]
260 for attr in attributes:
261 class_node.locals[attr] = fake.body[0].locals[attr]
262 # we use UseInferenceDefault, we can't be a generator so return an iterator
263 return iter([class_node])
264
265
266def _get_renamed_namedtuple_attributes(field_names):
267 names = list(field_names)
268 seen = set()
269 for i, name in enumerate(field_names):
270 if (
271 not all(c.isalnum() or c == "_" for c in name)
272 or keyword.iskeyword(name)
273 or not name
274 or name[0].isdigit()
275 or name.startswith("_")
276 or name in seen
277 ):
278 names[i] = "_%d" % i
279 seen.add(name)
280 return tuple(names)
281
282
283def _check_namedtuple_attributes(typename, attributes, rename=False):
284 attributes = tuple(attributes)
285 if rename:
286 attributes = _get_renamed_namedtuple_attributes(attributes)
287
288 # The following snippet is derived from the CPython Lib/collections/__init__.py sources
289 # <snippet>
290 for name in (typename, *attributes):
291 if not isinstance(name, str):
292 raise AstroidTypeError("Type names and field names must be strings")
293 if not name.isidentifier():
294 raise AstroidValueError(
295 "Type names and field names must be valid" + f"identifiers: {name!r}"
296 )
297 if keyword.iskeyword(name):
298 raise AstroidValueError(
299 f"Type names and field names cannot be a keyword: {name!r}"
300 )
301
302 seen = set()
303 for name in attributes:
304 if name.startswith("_") and not rename:
305 raise AstroidValueError(
306 f"Field names cannot start with an underscore: {name!r}"
307 )
308 if name in seen:
309 raise AstroidValueError(f"Encountered duplicate field name: {name!r}")
310 seen.add(name)
311 # </snippet>
312
313 return attributes
314
315
316def infer_enum(
317 node: nodes.Call, context: InferenceContext | None = None
318) -> Iterator[bases.Instance]:
319 """Specific inference function for enum Call node."""
320 # Raise `UseInferenceDefault` if `node` is a call to a a user-defined Enum.
321 try:
322 inferred = node.func.infer(context)
323 except (InferenceError, StopIteration) as exc:
324 raise UseInferenceDefault from exc
325
326 if not any(
327 isinstance(item, nodes.ClassDef) and item.qname() == ENUM_QNAME
328 for item in inferred
329 ):
330 raise UseInferenceDefault
331
332 enum_meta = _extract_single_node(
333 """
334 class EnumMeta(object):
335 'docstring'
336 def __call__(self, node):
337 class EnumAttribute(object):
338 name = ''
339 value = 0
340 return EnumAttribute()
341 def __iter__(self):
342 class EnumAttribute(object):
343 name = ''
344 value = 0
345 return [EnumAttribute()]
346 def __reversed__(self):
347 class EnumAttribute(object):
348 name = ''
349 value = 0
350 return (EnumAttribute, )
351 def __next__(self):
352 return next(iter(self))
353 def __getitem__(self, attr):
354 class Value(object):
355 @property
356 def name(self):
357 return ''
358 @property
359 def value(self):
360 return attr
361
362 return Value()
363 __members__ = ['']
364 """
365 )
366 class_node = infer_func_form(node, [enum_meta], context=context, enum=True)[0]
367 return iter([class_node.instantiate_class()])
368
369
370INT_FLAG_ADDITION_METHODS = """
371 def __or__(self, other):
372 return {name}(self.value | other.value)
373 def __and__(self, other):
374 return {name}(self.value & other.value)
375 def __xor__(self, other):
376 return {name}(self.value ^ other.value)
377 def __add__(self, other):
378 return {name}(self.value + other.value)
379 def __div__(self, other):
380 return {name}(self.value / other.value)
381 def __invert__(self):
382 return {name}(~self.value)
383 def __mul__(self, other):
384 return {name}(self.value * other.value)
385"""
386
387
388def infer_enum_class(node: nodes.ClassDef) -> nodes.ClassDef:
389 """Specific inference for enums."""
390 for basename in (b for cls in node.mro() for b in cls.basenames):
391 if node.root().name == "enum":
392 # Skip if the class is directly from enum module.
393 break
394 dunder_members = {}
395 target_names = set()
396 for local, values in node.locals.items():
397 if (
398 any(not isinstance(value, nodes.AssignName) for value in values)
399 or local == "_ignore_"
400 ):
401 continue
402
403 stmt = values[0].statement()
404 if isinstance(stmt, nodes.Assign):
405 if isinstance(stmt.targets[0], nodes.Tuple):
406 targets = stmt.targets[0].itered()
407 else:
408 targets = stmt.targets
409 elif isinstance(stmt, nodes.AnnAssign):
410 targets = [stmt.target]
411 else:
412 continue
413
414 inferred_return_value = None
415 if stmt.value is not None:
416 if isinstance(stmt.value, nodes.Const):
417 if isinstance(stmt.value.value, str):
418 inferred_return_value = repr(stmt.value.value)
419 else:
420 inferred_return_value = stmt.value.value
421 else:
422 inferred_return_value = stmt.value.as_string()
423
424 new_targets = []
425 for target in targets:
426 if isinstance(target, nodes.Starred):
427 continue
428 target_names.add(target.name)
429 # Replace all the assignments with our mocked class.
430 classdef = dedent(
431 """
432 class {name}({types}):
433 @property
434 def value(self):
435 return {return_value}
436 @property
437 def _value_(self):
438 return {return_value}
439 @property
440 def name(self):
441 return "{name}"
442 @property
443 def _name_(self):
444 return "{name}"
445 """.format(
446 name=target.name,
447 types=", ".join(node.basenames),
448 return_value=inferred_return_value,
449 )
450 )
451 if "IntFlag" in basename:
452 # Alright, we need to add some additional methods.
453 # Unfortunately we still can't infer the resulting objects as
454 # Enum members, but once we'll be able to do that, the following
455 # should result in some nice symbolic execution
456 classdef += INT_FLAG_ADDITION_METHODS.format(name=target.name)
457
458 fake = AstroidBuilder(
459 AstroidManager(), apply_transforms=False
460 ).string_build(classdef)[target.name]
461 fake.parent = target.parent
462 for method in node.mymethods():
463 fake.locals[method.name] = [method]
464 new_targets.append(fake.instantiate_class())
465 if stmt.value is None:
466 continue
467 dunder_members[local] = fake
468 node.locals[local] = new_targets
469
470 # The undocumented `_value2member_map_` member:
471 node.locals["_value2member_map_"] = [
472 nodes.Dict(
473 parent=node,
474 lineno=node.lineno,
475 col_offset=node.col_offset,
476 end_lineno=node.end_lineno,
477 end_col_offset=node.end_col_offset,
478 )
479 ]
480
481 members = nodes.Dict(
482 parent=node,
483 lineno=node.lineno,
484 col_offset=node.col_offset,
485 end_lineno=node.end_lineno,
486 end_col_offset=node.end_col_offset,
487 )
488 members.postinit(
489 [
490 (
491 nodes.Const(k, parent=members),
492 nodes.Name(
493 v.name,
494 parent=members,
495 lineno=v.lineno,
496 col_offset=v.col_offset,
497 end_lineno=v.end_lineno,
498 end_col_offset=v.end_col_offset,
499 ),
500 )
501 for k, v in dunder_members.items()
502 ]
503 )
504 node.locals["__members__"] = [members]
505 # The enum.Enum class itself defines two @DynamicClassAttribute data-descriptors
506 # "name" and "value" (which we override in the mocked class for each enum member
507 # above). When dealing with inference of an arbitrary instance of the enum
508 # class, e.g. in a method defined in the class body like:
509 # class SomeEnum(enum.Enum):
510 # def method(self):
511 # self.name # <- here
512 # In the absence of an enum member called "name" or "value", these attributes
513 # should resolve to the descriptor on that particular instance, i.e. enum member.
514 # For "value", we have no idea what that should be, but for "name", we at least
515 # know that it should be a string, so infer that as a guess.
516 if "name" not in target_names:
517 code = dedent(
518 """
519 @property
520 def name(self):
521 return ''
522 """
523 )
524 name_dynamicclassattr = AstroidBuilder(AstroidManager()).string_build(code)[
525 "name"
526 ]
527 node.locals["name"] = [name_dynamicclassattr]
528 break
529 return node
530
531
532def infer_typing_namedtuple_class(class_node, context: InferenceContext | None = None):
533 """Infer a subclass of typing.NamedTuple."""
534 # Check if it has the corresponding bases
535 annassigns_fields = [
536 annassign.target.name
537 for annassign in class_node.body
538 if isinstance(annassign, nodes.AnnAssign)
539 ]
540 code = dedent(
541 """
542 from collections import namedtuple
543 namedtuple({typename!r}, {fields!r})
544 """
545 ).format(typename=class_node.name, fields=",".join(annassigns_fields))
546 node = extract_node(code)
547 try:
548 generated_class_node = next(infer_named_tuple(node, context))
549 except StopIteration as e:
550 raise InferenceError(node=node, context=context) from e
551 for method in class_node.mymethods():
552 generated_class_node.locals[method.name] = [method]
553
554 for body_node in class_node.body:
555 if isinstance(body_node, nodes.Assign):
556 for target in body_node.targets:
557 attr = target.name
558 generated_class_node.locals[attr] = class_node.locals[attr]
559 elif isinstance(body_node, nodes.ClassDef):
560 generated_class_node.locals[body_node.name] = [body_node]
561
562 return iter((generated_class_node,))
563
564
565def infer_typing_namedtuple_function(node, context: InferenceContext | None = None):
566 """
567 Starting with python3.9, NamedTuple is a function of the typing module.
568 The class NamedTuple is build dynamically through a call to `type` during
569 initialization of the `_NamedTuple` variable.
570 """
571 klass = extract_node(
572 """
573 from typing import _NamedTuple
574 _NamedTuple
575 """
576 )
577 return klass.infer(context)
578
579
580def infer_typing_namedtuple(
581 node: nodes.Call, context: InferenceContext | None = None
582) -> Iterator[nodes.ClassDef]:
583 """Infer a typing.NamedTuple(...) call."""
584 # This is essentially a namedtuple with different arguments
585 # so we extract the args and infer a named tuple.
586 try:
587 func = next(node.func.infer())
588 except (InferenceError, StopIteration) as exc:
589 raise UseInferenceDefault from exc
590
591 if func.qname() not in TYPING_NAMEDTUPLE_QUALIFIED:
592 raise UseInferenceDefault
593
594 if len(node.args) != 2:
595 raise UseInferenceDefault
596
597 if not isinstance(node.args[1], (nodes.List, nodes.Tuple)):
598 raise UseInferenceDefault
599
600 return infer_named_tuple(node, context)
601
602
603def _get_namedtuple_fields(node: nodes.Call) -> str:
604 """Get and return fields of a NamedTuple in code-as-a-string.
605
606 Because the fields are represented in their code form we can
607 extract a node from them later on.
608 """
609 names = []
610 container = None
611 try:
612 container = next(node.args[1].infer())
613 except (InferenceError, StopIteration) as exc:
614 raise UseInferenceDefault from exc
615 # We pass on IndexError as we'll try to infer 'field_names' from the keywords
616 except IndexError:
617 pass
618 if not container:
619 for keyword_node in node.keywords:
620 if keyword_node.arg == "field_names":
621 try:
622 container = next(keyword_node.value.infer())
623 except (InferenceError, StopIteration) as exc:
624 raise UseInferenceDefault from exc
625 break
626 if not isinstance(container, nodes.BaseContainer):
627 raise UseInferenceDefault
628 for elt in container.elts:
629 if isinstance(elt, nodes.Const):
630 names.append(elt.as_string())
631 continue
632 if not isinstance(elt, (nodes.List, nodes.Tuple)):
633 raise UseInferenceDefault
634 if len(elt.elts) != 2:
635 raise UseInferenceDefault
636 names.append(elt.elts[0].as_string())
637
638 if names:
639 field_names = f"({','.join(names)},)"
640 else:
641 field_names = ""
642 return field_names
643
644
645def _is_enum_subclass(cls: astroid.ClassDef) -> bool:
646 """Return whether cls is a subclass of an Enum."""
647 return cls.is_subtype_of("enum.Enum")
648
649
650def register(manager: AstroidManager) -> None:
651 manager.register_transform(
652 nodes.Call, inference_tip(infer_named_tuple), _looks_like_namedtuple
653 )
654 manager.register_transform(nodes.Call, inference_tip(infer_enum), _looks_like_enum)
655 manager.register_transform(
656 nodes.ClassDef, infer_enum_class, predicate=_is_enum_subclass
657 )
658 manager.register_transform(
659 nodes.ClassDef,
660 inference_tip(infer_typing_namedtuple_class),
661 _has_namedtuple_base,
662 )
663 manager.register_transform(
664 nodes.FunctionDef,
665 inference_tip(infer_typing_namedtuple_function),
666 lambda node: node.name == "NamedTuple"
667 and getattr(node.root(), "name", None) == "typing",
668 )
669 manager.register_transform(
670 nodes.Call,
671 inference_tip(infer_typing_namedtuple),
672 _looks_like_typing_namedtuple,
673 )