Allow nested classes in NamedTuple bodies

python · sobolevn · Jul 29, 2023 · Jul 29, 2023 · Jul 29, 2023 · Jul 29, 2023
commit 9150e8b9ff0789864e6afd4fc2693222efdbc7e4
diff --git a/mypy/semanal_namedtuple.py b/mypy/semanal_namedtuple.py
@@ -149,7 +149,7 @@ def check_namedtuple_classdef(
         default_items: dict[str, Expression] = {}
         statements: list[Statement] = []
         for stmt in defn.defs.body:
-            statements.append(stmt)
+            # Processing fields of a namedtuple:
             if not isinstance(stmt, AssignmentStmt):
                 # Still allow pass or ... (for empty namedtuples).
                 if isinstance(stmt, PassStmt) or (
@@ -162,16 +162,20 @@ def check_namedtuple_classdef(
                 # And docstrings.
                 if isinstance(stmt, ExpressionStmt) and isinstance(stmt.expr, StrExpr):
                     continue
-                statements.pop()
+                # And nested classes, they need to be analyzed further:
+                if isinstance(stmt, ClassDef):
+                    statements.append(stmt)
+                    continue
+
                 defn.removed_statements.append(stmt)
                 self.fail(NAMEDTUP_CLASS_ERROR, stmt)
             elif len(stmt.lvalues) > 1 or not isinstance(stmt.lvalues[0], NameExpr):
                 # An assignment, but an invalid one.
-                statements.pop()
                 defn.removed_statements.append(stmt)
                 self.fail(NAMEDTUP_CLASS_ERROR, stmt)
             else:
                 # Append name and type in this case...
+                statements.append(stmt)
                 name = stmt.lvalues[0].name
                 items.append(name)
                 if stmt.type is None:

diff --git a/test-data/unit/check-class-namedtuple.test b/test-data/unit/check-class-namedtuple.test
@@ -368,6 +368,40 @@ class X(NamedTuple):
     y = 2  # E: Invalid statement in NamedTuple definition; expected "field_name: field_type [= default]"
 [builtins fixtures/tuple.pyi]
 
+[case testNewNamedTupleWithNestedClass]
+from typing import NamedTuple
+
+class A(NamedTuple):
+    x: int
+    class B:
+        x: str
+
+a: A
+reveal_type(A.B)  # N: Revealed type is "def () -> __main__.A.B"
+b: A.B
+reveal_type(b.x)  # N: Revealed type is "builtins.str"
+[builtins fixtures/tuple.pyi]
+
+[case testNewNamedTupleWithNestedNamedTuple]
+from typing import NamedTuple
+
+class A(NamedTuple):
+    x: int
+    class B(NamedTuple):
+        x: str
+        y: int = 1
+
+# Correct:
+A(1)
+A.B('a')
+A.B('a', 2)
+
+# Incorrect:
+A.B()  # E: Missing positional argument "x" in call to "B"
+A.B(1, 'a')  # E: Argument 1 to "B" has incompatible type "int"; expected "str" \
+             # E: Argument 2 to "B" has incompatible type "str"; expected "int"
+[builtins fixtures/tuple.pyi]
+
 [case testTypeUsingTypeCNamedTuple]
 from typing import NamedTuple, Type