Avoid PEP 604 upgrades that lead to invalid syntax (astral-sh#6888)

Closes astral-sh#6843.

crates/ruff/resources/test/fixtures/pyupgrade/UP007.py

@@ -59,3 +59,35 @@ def f() -> None:
     x = Union["str", "int"]
     x: Union[str, int]
     x: Union["str", "int"]
+
+
+def f(x: Union[int : float]) -> None:
+    ...
+
+
+def f(x: Union[str, int : float]) -> None:
+    ...
+
+
+def f(x: Union[x := int]) -> None:
+    ...
+
+
+def f(x: Union[str, x := int]) -> None:
+    ...
+
+
+def f(x: Union[lambda: int]) -> None:
+    ...
+
+
+def f(x: Union[str, lambda: int]) -> None:
+    ...
+
+
+def f(x: Optional[int : float]) -> None:
+    ...
+
+
+def f(x: Optional[str, int : float]) -> None:
+    ...

crates/ruff/src/rules/pyupgrade/rules/use_pep604_annotation.rs

@@ -66,9 +66,11 @@ pub(crate) fn use_pep604_annotation(
     slice: &Expr,
     operator: Pep604Operator,
 ) {
-    // Avoid fixing forward references, or types not in an annotation.
+    // Avoid fixing forward references, types not in an annotation, and expressions that would
+    // lead to invalid syntax.
     let fixable = checker.semantic().in_type_definition()
-        && !checker.semantic().in_complex_string_type_definition();
+        && !checker.semantic().in_complex_string_type_definition()
+        && is_allowed_value(slice);
 
     match operator {
         Pep604Operator::Optional => {
@@ -128,3 +130,52 @@ fn union(elts: &[Expr], locator: &Locator) -> String {
         elts.map(|expr| locator.slice(expr.range())).join(" | ")
     }
 }
+
+/// Returns `true` if the expression is valid for use in a bitwise union (e.g., `X | Y`). Returns
+/// `false` for lambdas, yield expressions, and other expressions that are invalid in such a
+/// context.
+fn is_allowed_value(expr: &Expr) -> bool {
+    // TODO(charlie): If the expression requires parentheses when multi-line, and the annotation
+    // itself is not parenthesized, this should return `false`. Consider, for example:
+    // ```python
+    // x: Union[
+    //     "Sequence["
+    //         "int"
+    //     "]",
+    //     float,
+    // ]
+    // ```
+    // Converting this to PEP 604 syntax requires that the multiline string is parenthesized.
+    match expr {
+        Expr::BoolOp(_)
+        | Expr::BinOp(_)
+        | Expr::UnaryOp(_)
+        | Expr::IfExp(_)
+        | Expr::Dict(_)
+        | Expr::Set(_)
+        | Expr::ListComp(_)
+        | Expr::SetComp(_)
+        | Expr::DictComp(_)
+        | Expr::GeneratorExp(_)
+        | Expr::Compare(_)
+        | Expr::Call(_)
+        | Expr::FormattedValue(_)
+        | Expr::FString(_)
+        | Expr::Constant(_)
+        | Expr::Attribute(_)
+        | Expr::Subscript(_)
+        | Expr::Name(_)
+        | Expr::List(_) => true,
+        Expr::Tuple(tuple) => tuple.elts.iter().all(is_allowed_value),
+        // Maybe require parentheses.
+        Expr::NamedExpr(_) => false,
+        // Invalid in binary expressions.
+        Expr::Await(_)
+        | Expr::Lambda(_)
+        | Expr::Yield(_)
+        | Expr::YieldFrom(_)
+        | Expr::Starred(_)
+        | Expr::Slice(_)
+        | Expr::IpyEscapeCommand(_) => false,
+    }
+}

crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP007.py.snap

@@ -275,6 +275,8 @@ UP007.py:60:8: UP007 [*] Use `X | Y` for type annotations
 60    |-    x: Union[str, int]
    60 |+    x: str | int
 61 61 |     x: Union["str", "int"]
+62 62 | 
+63 63 | 
 
 UP007.py:61:8: UP007 [*] Use `X | Y` for type annotations
    |
@@ -291,5 +293,72 @@ UP007.py:61:8: UP007 [*] Use `X | Y` for type annotations
 60 60 |     x: Union[str, int]
 61    |-    x: Union["str", "int"]
    61 |+    x: "str" | "int"
+62 62 | 
+63 63 | 
+64 64 | def f(x: Union[int : float]) -> None:
+
+UP007.py:64:10: UP007 Use `X | Y` for type annotations
+   |
+64 | def f(x: Union[int : float]) -> None:
+   |          ^^^^^^^^^^^^^^^^^^ UP007
+65 |     ...
+   |
+   = help: Convert to `X | Y`
+
+UP007.py:68:10: UP007 Use `X | Y` for type annotations
+   |
+68 | def f(x: Union[str, int : float]) -> None:
+   |          ^^^^^^^^^^^^^^^^^^^^^^^ UP007
+69 |     ...
+   |
+   = help: Convert to `X | Y`
+
+UP007.py:72:10: UP007 Use `X | Y` for type annotations
+   |
+72 | def f(x: Union[x := int]) -> None:
+   |          ^^^^^^^^^^^^^^^ UP007
+73 |     ...
+   |
+   = help: Convert to `X | Y`
+
+UP007.py:76:10: UP007 Use `X | Y` for type annotations
+   |
+76 | def f(x: Union[str, x := int]) -> None:
+   |          ^^^^^^^^^^^^^^^^^^^^ UP007
+77 |     ...
+   |
+   = help: Convert to `X | Y`
+
+UP007.py:80:10: UP007 Use `X | Y` for type annotations
+   |
+80 | def f(x: Union[lambda: int]) -> None:
+   |          ^^^^^^^^^^^^^^^^^^ UP007
+81 |     ...
+   |
+   = help: Convert to `X | Y`
+
+UP007.py:84:10: UP007 Use `X | Y` for type annotations
+   |
+84 | def f(x: Union[str, lambda: int]) -> None:
+   |          ^^^^^^^^^^^^^^^^^^^^^^^ UP007
+85 |     ...
+   |
+   = help: Convert to `X | Y`
+
+UP007.py:88:10: UP007 Use `X | Y` for type annotations
+   |
+88 | def f(x: Optional[int : float]) -> None:
+   |          ^^^^^^^^^^^^^^^^^^^^^ UP007
+89 |     ...
+   |
+   = help: Convert to `X | Y`
+
+UP007.py:92:10: UP007 Use `X | Y` for type annotations
+   |
+92 | def f(x: Optional[str, int : float]) -> None:
+   |          ^^^^^^^^^^^^^^^^^^^^^^^^^^ UP007
+93 |     ...
+   |
+   = help: Convert to `X | Y`