From 2b51a9b41b902238cfe061cc91d03b23378011b6 Mon Sep 17 00:00:00 2001
From: Ali Hamdan <ali.hamdan.dev@gmail.com>
Date: Fri, 3 Feb 2023 22:56:39 +0100
Subject: [PATCH] [stubgen] Preserve PEP 604 Unions in generated pyi files

When a PEP 604 Union exists in the runtime, stubgen was generating a
`Union[...]` syntax without importing `Union` from `typing`. With this
change, stubgen preserves the ` | `-unions in the output.

Fixes #12929
Closes #13428
Ref #12920
---
 mypy/stubgen.py             |  4 ++++
 test-data/unit/stubgen.test | 10 ++++++++++
 2 files changed, 14 insertions(+)

diff --git a/mypy/stubgen.py b/mypy/stubgen.py
index 51ee1b93de14..6cb4669887fe 100755
--- a/mypy/stubgen.py
+++ b/mypy/stubgen.py
@@ -134,6 +134,7 @@
     TypeList,
     TypeStrVisitor,
     UnboundType,
+    UnionType,
     get_proper_type,
 )
 from mypy.visitor import NodeVisitor
@@ -326,6 +327,9 @@ def visit_none_type(self, t: NoneType) -> str:
     def visit_type_list(self, t: TypeList) -> str:
         return f"[{self.list_str(t.items)}]"
 
+    def visit_union_type(self, t: UnionType) -> str:
+        return " | ".join([item.accept(self) for item in t.items])
+
     def args_str(self, args: Iterable[Type]) -> str:
         """Convert an array of arguments to strings and join the results with commas.
 
diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test
index 4909c0005412..8e4285b7de2e 100644
--- a/test-data/unit/stubgen.test
+++ b/test-data/unit/stubgen.test
@@ -2783,3 +2783,13 @@ T = TypeVar("T", bound=str | None)
 from typing import TypeVar
 
 T = TypeVar('T', bound=str | None)
+
+
+[case testPEP604UnionType]
+a: str | int
+
+def f(x: str | None) -> None: ...
+[out]
+a: str | int
+
+def f(x: str | None) -> None: ...