Process NameTuple and fix Union types using |

mypy/stubgen.py

@@ -48,7 +48,7 @@
 import sys
 import traceback
 from collections import defaultdict
-from typing import Iterable, List, Mapping, cast
+from typing import Iterable, List, Mapping
 from typing_extensions import Final
 
 import mypy.build
@@ -94,6 +94,7 @@
     ListExpr,
     MemberExpr,
     MypyFile,
+    NamedTupleExpr,
     NameExpr,
     OverloadedFuncDef,
     Statement,
@@ -133,6 +134,7 @@
     TypeList,
     TypeStrVisitor,
     UnboundType,
+    UnionType,
     get_proper_type,
 )
 from mypy.visitor import NodeVisitor
@@ -303,6 +305,11 @@ def visit_unbound_type(self, t: UnboundType) -> str:
             s += f"[{self.args_str(t.args)}]"
         return s
 
+    def visit_union_type(self, t: UnionType) -> str:
+        s = super().visit_union_type(t)
+        self.stubgen.import_tracker.require_name("Union")
+        return s
+
     def visit_none_type(self, t: NoneType) -> str:
         return "None"
 
@@ -599,6 +606,7 @@ def __init__(
         self.export_less = export_less
         # Add imports that could be implicitly generated
         self.import_tracker.add_import_from("typing", [("NamedTuple", None)])
+        self.import_tracker.add_import_from("typing", [("Union", None)])
         # Names in __all__ are required
         for name in _all_ or ():
             if name not in IGNORED_DUNDERS:
@@ -932,6 +940,16 @@ def visit_class_def(self, o: ClassDef) -> None:
             self.add(f"({', '.join(base_types)})")
         self.add(":\n")
         n = len(self._output)
+
+        if "namedtuple" in base_types or "NamedTuple" in base_types:
+            for expr in o.base_type_exprs:
+                items: list[tuple[str | None, str | None]] = []
+                if isinstance(expr, CallExpr):
+                    parameters = self.process_namedtuple_parameters(expr.args[1])
+                    if parameters is not None:
+                        items.extend(parameters)
+                self.print_named_tuple(items)
+
         self._indent += "    "
         self._vars.append([])
         super().visit_class_def(o)
@@ -960,6 +978,9 @@ def get_base_types(self, cdef: ClassDef) -> list[str]:
             elif isinstance(base, IndexExpr):
                 p = AliasPrinter(self)
                 base_types.append(base.accept(p))
+            elif isinstance(base, CallExpr) and isinstance(base.analyzed, NamedTupleExpr):
+                base_types.append("NamedTuple")
+
         return base_types
 
     def visit_block(self, o: Block) -> None:
@@ -1017,33 +1038,61 @@ def is_namedtuple(self, expr: Expression) -> bool:
         if not isinstance(expr, CallExpr):
             return False
         callee = expr.callee
-        return (isinstance(callee, NameExpr) and callee.name.endswith("namedtuple")) or (
-            isinstance(callee, MemberExpr) and callee.name == "namedtuple"
+        return (
+            isinstance(callee, NameExpr)
+            and (callee.name.endswith("namedtuple") or callee.name.endswith("NamedTuple"))
+        ) or (
+            isinstance(callee, MemberExpr)
+            and (callee.name == "namedtuple" or callee.name == "NamedTuple")
         )
 
+    def process_namedtuple_parameters(
+        self, rvalue: Expression
+    ) -> list[tuple[str | None, str | None]] | None:
+        if isinstance(rvalue, StrExpr):
+            items: list[tuple[str | None, str | None]] = [
+                (key, "Incomplete") for key in rvalue.value.replace(",", " ").split()
+            ]
+        elif isinstance(rvalue, (ListExpr, TupleExpr)):
+            items = [self.process_namedtuple_type(item) for item in rvalue.items]
+        else:
+            return None
+        return list(filter(lambda v: v[0] is not None and v[0].isidentifier(), items))
+
     def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
         if self._state != EMPTY:
             self.add("\n")
-        if isinstance(rvalue.args[1], StrExpr):
-            items = rvalue.args[1].value.replace(",", " ").split()
-        elif isinstance(rvalue.args[1], (ListExpr, TupleExpr)):
-            list_items = cast(List[StrExpr], rvalue.args[1].items)
-            items = [item.value for item in list_items]
-        else:
+
+        items = self.process_namedtuple_parameters(rvalue.args[1])
+        if items is None:
             self.add(f"{self._indent}{lvalue.name}: Incomplete")
             self.import_tracker.require_name("Incomplete")
             return
+
         self.import_tracker.require_name("NamedTuple")
         self.add(f"{self._indent}class {lvalue.name}(NamedTuple):")
+
         if len(items) == 0:
             self.add(" ...\n")
         else:
-            self.import_tracker.require_name("Incomplete")
             self.add("\n")
-            for item in items:
-                self.add(f"{self._indent}    {item}: Incomplete\n")
+            self.print_named_tuple(items)
         self._state = CLASS
 
+    def print_named_tuple(self, items: list[tuple[str | None, str | None]]) -> None:
+        self.import_tracker.require_name("Incomplete")
+        for item in items:
+            key, rtype = item
+            self.add(f"{self._indent}    {key}: {rtype}\n")
+
+    def process_namedtuple_type(self, item: Expression) -> tuple[str | None, str | None]:
+        if isinstance(item, StrExpr):
+            return item.value, "Incomplete"
+        elif isinstance(item, TupleExpr) and isinstance(item.items[0], StrExpr):
+            p = AliasPrinter(self)
+            return item.items[0].value, item.items[1].accept(p)
+        return None, None
+
     def is_alias_expression(self, expr: Expression, top_level: bool = True) -> bool:
         """Return True for things that look like target for an alias.
 

test-data/unit/stubgen.test

@@ -516,6 +516,17 @@ __all__ = ['urllib']
 [out]
 import urllib as urllib
 
+
+[case testImportUnion]
+def func(a: str | int) -> str | int:
+  pass
+
+[out]
+from typing import Union
+
+def func(a: Union[str, int]) -> Union[str, int]: ...
+
+
 [case testRelativeImportAll]
 from .x import *
 [out]
@@ -595,6 +606,33 @@ class X(NamedTuple):
     a: Incomplete
     b: Incomplete
 
+[case testNamedtupleUsingInvalidIdent]
+import collections, x
+X = collections.namedtuple('X', ['@'])
+[out]
+from typing import NamedTuple
+
+class X(NamedTuple): ...
+
+[case testNamedtupleWithTypesInvalidIdent]
+import collections, x
+X = typing.NamedTuple('X', [('@', int)])
+[out]
+from typing import NamedTuple
+
+class X(NamedTuple): ...
+
+[case testNamedtupleWithTypes]
+import collections, x
+X = typing.NamedTuple('X', [('a', str), ('b', str)])
+[out]
+from _typeshed import Incomplete
+from typing import NamedTuple
+
+class X(NamedTuple):
+    a: str
+    b: str
+
 [case testEmptyNamedtuple]
 import collections
 X = collections.namedtuple('X', [])
@@ -915,7 +953,7 @@ T = TypeVar('T')
 alias = Union[T, List[T]]
 
 [out]
-from typing import TypeVar
+from typing import TypeVar, Union
 
 T = TypeVar('T')
 alias = Union[T, List[T]]
@@ -2705,3 +2743,19 @@ def f():
     return 0
 [out]
 def f(): ...
+
+[case testClassNamedtuple_semanal]
+import typing
+import collections, x
+class X(typing.NamedTuple("X", [("a", int), ("b", int)])):
+    def m(self):
+       pass
+
+[out]
+from _typeshed import Incomplete
+from typing import NamedTuple
+
+class X(NamedTuple):
+    a: int
+    b: int
+    def m(self) -> None: ...