From 917cc75fd6f1417edb45eb77e449934f794c18fc Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Sun, 11 Aug 2024 22:26:32 +0100
Subject: [PATCH] An alternative fix for a union-like literal string (#17639)

It is unfortunate to add two extra slots to a common type (and I guess
this is why it was rejected in the original PR), but all other
alternatives I tried are hacky and/or dangerous. So, this is a price to
pay for introducing a new type syntax.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
(cherry picked from commit 3da16bd6004ec9685c476f77526abd36e7632813)
---
 mypy/fastparse.py                 |  4 +---
 mypy/typeanal.py                  |  6 +++++-
 mypy/types.py                     | 16 +++++++++++++---
 test-data/unit/check-literal.test |  4 ++++
 4 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/mypy/fastparse.py b/mypy/fastparse.py
index 0d93c368585d..c78a24ef2a15 100644
--- a/mypy/fastparse.py
+++ b/mypy/fastparse.py
@@ -329,12 +329,10 @@ def parse_type_string(
     """
     try:
         _, node = parse_type_comment(f"({expr_string})", line=line, column=column, errors=None)
-        if isinstance(node, UnboundType) and node.original_str_expr is None:
+        if isinstance(node, (UnboundType, UnionType)) and node.original_str_expr is None:
             node.original_str_expr = expr_string
             node.original_str_fallback = expr_fallback_name
             return node
-        elif isinstance(node, UnionType):
-            return node
         else:
             return RawExpressionType(expr_string, expr_fallback_name, line, column)
     except (SyntaxError, ValueError):
diff --git a/mypy/typeanal.py b/mypy/typeanal.py
index 5613e01a272b..ee6cd8765aae 100644
--- a/mypy/typeanal.py
+++ b/mypy/typeanal.py
@@ -1575,7 +1575,11 @@ def analyze_literal_type(self, t: UnboundType) -> Type:
 
     def analyze_literal_param(self, idx: int, arg: Type, ctx: Context) -> list[Type] | None:
         # This UnboundType was originally defined as a string.
-        if isinstance(arg, UnboundType) and arg.original_str_expr is not None:
+        if (
+            isinstance(arg, ProperType)
+            and isinstance(arg, (UnboundType, UnionType))
+            and arg.original_str_expr is not None
+        ):
             assert arg.original_str_fallback is not None
             return [
                 LiteralType(
diff --git a/mypy/types.py b/mypy/types.py
index 1fcd5b920ab6..6db68629ef6f 100644
--- a/mypy/types.py
+++ b/mypy/types.py
@@ -914,7 +914,7 @@ class UnboundType(ProperType):
 
     def __init__(
         self,
-        name: str | None,
+        name: str,
         args: Sequence[Type] | None = None,
         line: int = -1,
         column: int = -1,
@@ -926,7 +926,6 @@ def __init__(
         super().__init__(line, column)
         if not args:
             args = []
-        assert name is not None
         self.name = name
         self.args = tuple(args)
         # Should this type be wrapped in an Optional?
@@ -2847,7 +2846,13 @@ def is_singleton_type(self) -> bool:
 class UnionType(ProperType):
     """The union type Union[T1, ..., Tn] (at least one type argument)."""
 
-    __slots__ = ("items", "is_evaluated", "uses_pep604_syntax")
+    __slots__ = (
+        "items",
+        "is_evaluated",
+        "uses_pep604_syntax",
+        "original_str_expr",
+        "original_str_fallback",
+    )
 
     def __init__(
         self,
@@ -2866,6 +2871,11 @@ def __init__(
         self.is_evaluated = is_evaluated
         # uses_pep604_syntax is True if Union uses OR syntax (X | Y)
         self.uses_pep604_syntax = uses_pep604_syntax
+        # The meaning of these two is the same as for UnboundType. A UnionType can be
+        # return by type parser from a string "A|B", and we need to be able to fall back
+        # to plain string, when such a string appears inside a Literal[...].
+        self.original_str_expr: str | None = None
+        self.original_str_fallback: str | None = None
 
     def can_be_true_default(self) -> bool:
         return any(item.can_be_true for item in self.items)
diff --git a/test-data/unit/check-literal.test b/test-data/unit/check-literal.test
index 2f5fa1b2ede7..8f8aaf6a3982 100644
--- a/test-data/unit/check-literal.test
+++ b/test-data/unit/check-literal.test
@@ -12,8 +12,12 @@ reveal_type(g1)  # N: Revealed type is "def (x: Literal['A['])"
 
 def f2(x: 'A B') -> None: pass  # E: Invalid type comment or annotation
 def g2(x: Literal['A B']) -> None: pass
+def h2(x: 'A|int') -> None: pass  # E: Name "A" is not defined
+def i2(x: Literal['A|B']) -> None: pass
 reveal_type(f2)  # N: Revealed type is "def (x: Any)"
 reveal_type(g2)  # N: Revealed type is "def (x: Literal['A B'])"
+reveal_type(h2)  # N: Revealed type is "def (x: Union[Any, builtins.int])"
+reveal_type(i2)  # N: Revealed type is "def (x: Literal['A|B'])"
 [builtins fixtures/tuple.pyi]
 [out]