Skip to content

Commit

Permalink
Refactor field inference to _get_namedtuple_fields (#1413)
Browse files Browse the repository at this point in the history
  • Loading branch information
DanielNoord committed Jul 1, 2022
1 parent 49beab3 commit 9531e03
Showing 1 changed file with 23 additions and 9 deletions.
32 changes: 23 additions & 9 deletions astroid/brain/brain_namedtuple_enum.py
Expand Up @@ -85,13 +85,17 @@ def infer_func_form(
try:
name, names = _find_func_form_arguments(node, context)
try:
attributes = names.value.replace(",", " ").split()
attributes: list[str] = names.value.replace(",", " ").split()
except AttributeError as exc:
# Handle attributes of NamedTuples
if not enum:
attributes = [
_infer_first(const, context).value for const in names.elts
]
attributes = []
fields = _get_namedtuple_fields(node)
if fields:
fields_node = extract_node(fields)
attributes = [
_infer_first(const, context).value for const in fields_node.elts
]

# Handle attributes of Enums
else:
Expand Down Expand Up @@ -522,21 +526,31 @@ def infer_typing_namedtuple(
if not isinstance(node.args[1], (nodes.List, nodes.Tuple)):
raise UseInferenceDefault

return infer_named_tuple(node, context)


def _get_namedtuple_fields(node: nodes.Call) -> str:
"""Get and return fields of a NamedTuple in code-as-a-string.
Because the fields are represented in their code form we can
extract a node from them later on.
"""
names = []
for elt in node.args[1].elts:
for elt in next(node.args[1].infer()).elts:
if isinstance(elt, nodes.Const):
names.append(elt.as_string())
continue
if not isinstance(elt, (nodes.List, nodes.Tuple)):
raise UseInferenceDefault
if len(elt.elts) != 2:
raise UseInferenceDefault
names.append(elt.elts[0].as_string())

typename = node.args[0].as_string()
if names:
field_names = f"({','.join(names)},)"
else:
field_names = "''"
node = extract_node(f"namedtuple({typename}, {field_names})")
return infer_named_tuple(node, context)
field_names = ""
return field_names


def _is_enum_subclass(cls: astroid.ClassDef) -> bool:
Expand Down

0 comments on commit 9531e03

Please sign in to comment.