Skip to content

Commit

Permalink
Recognize nested classes in classes inherited from NamedTuple
Browse files Browse the repository at this point in the history
  • Loading branch information
dimp-gh authored and Pierre-Sassoulas committed Sep 6, 2021
1 parent 46628ba commit ac100bc
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 7 deletions.
14 changes: 7 additions & 7 deletions astroid/brain/brain_namedtuple_enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -484,13 +484,13 @@ def infer_typing_namedtuple_class(class_node, context=None):
for method in class_node.mymethods():
generated_class_node.locals[method.name] = [method]

for assign in class_node.body:
if not isinstance(assign, nodes.Assign):
continue

for target in assign.targets:
attr = target.name
generated_class_node.locals[attr] = class_node.locals[attr]
for body_node in class_node.body:
if isinstance(body_node, nodes.Assign):
for target in body_node.targets:
attr = target.name
generated_class_node.locals[attr] = class_node.locals[attr]
elif isinstance(body_node, nodes.ClassDef):
generated_class_node.locals[body_node.name] = [body_node]

return iter((generated_class_node,))

Expand Down
21 changes: 21 additions & 0 deletions tests/unittest_brain.py
Original file line number Diff line number Diff line change
Expand Up @@ -1629,6 +1629,27 @@ def test_typing_types(self):
inferred = next(node.infer())
self.assertIsInstance(inferred, nodes.ClassDef, node.as_string())

def test_namedtuple_nested_class(self):
result = builder.extract_node(
"""
from typing import NamedTuple
class Example(NamedTuple):
class Foo:
bar = "bar"
Example
"""
)
inferred = next(result.infer())
self.assertIsInstance(inferred, astroid.ClassDef)

class_def_attr = inferred.getattr("Foo")[0]
self.assertIsInstance(class_def_attr, astroid.ClassDef)
attr_def = class_def_attr.getattr("bar")[0]
attr = next(attr_def.infer())
self.assertEqual(attr.value, "bar")

@test_utils.require_version(minver="3.7")
def test_tuple_type(self):
node = builder.extract_node(
Expand Down

0 comments on commit ac100bc

Please sign in to comment.