An alternative fix for a union-like literal string (python#17639)

ilevkivskyi · pre-commit-ci[bot] · web-flow · commit 3da16bd6004e · 2024-08-11T22:26:32.000+01:00
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] &lt;66853113+pre-commit-ci[bot]@users.noreply.github.com&gt;
diff --git a/mypy/fastparse.py b/mypy/fastparse.py
@@ -331,12 +331,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
@@ -1615,7 +1615,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
@@ -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?
@@ -2849,7 +2848,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,
@@ -2868,6 +2873,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
@@ -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]
 
