Skip to content

Add some typing to NamedTuple brain #1412

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Feb 28, 2022
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions astroid/brain/brain_namedtuple_enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,12 @@
import functools
import keyword
from textwrap import dedent
from typing import Iterator, List, Optional, Tuple

import astroid
from astroid import arguments, inference_tip, nodes, util
from astroid.builder import AstroidBuilder, extract_node
from astroid.context import InferenceContext
from astroid.exceptions import (
AstroidTypeError,
AstroidValueError,
Expand Down Expand Up @@ -90,7 +92,12 @@ def _extract_namedtuple_arg_or_keyword( # pylint: disable=inconsistent-return-s
raise UseInferenceDefault()


def infer_func_form(node, base_type, context=None, enum=False):
def infer_func_form(
node: nodes.Call,
base_type: nodes.NodeNG,
context: Optional[InferenceContext] = None,
enum: bool = False,
) -> Tuple[nodes.ClassDef, str, List[str]]:
"""Specific inference function for namedtuple or Python 3 enum."""
# node is a Call node, class name as first argument and generated class
# attributes as second argument
Expand All @@ -102,10 +109,13 @@ def infer_func_form(node, base_type, context=None, enum=False):
try:
attributes = 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
]

# Handle attributes of Enums
else:
# Enums supports either iterator of (name, value) pairs
# or mappings.
Expand Down Expand Up @@ -183,7 +193,9 @@ def _looks_like(node, name):
_looks_like_typing_namedtuple = functools.partial(_looks_like, name="NamedTuple")


def infer_named_tuple(node, context=None):
def infer_named_tuple(
node: nodes.Call, context: Optional[InferenceContext] = None
) -> Iterator[nodes.ClassDef]:
"""Specific inference function for namedtuple Call node"""
tuple_base_name = nodes.Name(name="tuple", parent=node.root())
class_node, name, attributes = infer_func_form(
Expand Down Expand Up @@ -507,7 +519,9 @@ def infer_typing_namedtuple_function(node, context=None):
return klass.infer(context)


def infer_typing_namedtuple(node, context=None):
def infer_typing_namedtuple(
node: nodes.Call, context: Optional[InferenceContext] = None
) -> Iterator[nodes.ClassDef]:
"""Infer a typing.NamedTuple(...) call."""
# This is essentially a namedtuple with different arguments
# so we extract the args and infer a named tuple.
Expand Down