Process NameTuple and fix Union types using |

Author: Matej Spiller Muys

mypy/stubgen.py

@@ -133,6 +133,7 @@
     TypeList,
     TypeStrVisitor,
     UnboundType,
+    UnionType,
     get_proper_type,
 )
 from mypy.visitor import NodeVisitor
@@ -303,6 +304,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 +605,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:
@@ -1017,18 +1024,24 @@ 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(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()
+            items: list[tuple[str, str | None] | None] = [
+                (key, "Incomplete") for key in 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]
+            items = [self.process_namedtuple_type(item) for item in list_items]
         else:
             self.add(f"{self._indent}{lvalue.name}: Incomplete")
             self.import_tracker.require_name("Incomplete")
@@ -1041,9 +1054,20 @@ def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
             self.import_tracker.require_name("Incomplete")
             self.add("\n")
             for item in items:
-                self.add(f"{self._indent}    {item}: Incomplete\n")
+                if item is None:
+                    continue
+                key, rtype = item
+                self.add(f"{self._indent}    {key}: {rtype}\n")
         self._state = CLASS
 
+    def process_namedtuple_type(self, item: StrExpr | TupleExpr) -> tuple[str, str | None] | None:
+        if isinstance(item, StrExpr):
+            return item.value, "Incomplete"
+        elif isinstance(item.items[0], StrExpr):
+            p = AliasPrinter(self)
+            return item.items[0].value, item.items[1].accept(p)
+        return 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,17 @@ class X(NamedTuple):
     a: Incomplete
     b: Incomplete
 
+[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 +937,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]]