From 932c9a47893c5d599e605112c732be737c3bf6f7 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Thu, 13 Jul 2023 07:34:04 -0400
Subject: [PATCH] Extend PEP 604 rewrites to support some quoted annotations
 (#5725)

## Summary

Python doesn't allow `"Foo" | None` if the annotation will be evaluated
at runtime (see the comments in the PR, or the semantic model
documentation for more on what this means and when it is true), but it
_does_ allow it if the annotation is typing-only.

This, for example, is invalid, as Python will evaluate `"Foo" | None` at
runtime in order to
populate the function's `__annotations__`:

```python
def f(x: "Foo" | None): ...
```

This, however, is valid:

```python
def f():
    x: "Foo" | None
```

As is this:

```python
from __future__ import annotations

def f(x: "Foo" | None): ...
```

Closes #5706.
---
 .../pyupgrade/rules/use_pep604_annotation.rs  |  1 +
 ...ff__rules__pyupgrade__tests__UP007.py.snap | 16 +++++++++++
 .../src/analyze/typing.rs                     | 28 ++++++++++++++-----
 3 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/crates/ruff/src/rules/pyupgrade/rules/use_pep604_annotation.rs b/crates/ruff/src/rules/pyupgrade/rules/use_pep604_annotation.rs
index b1515b410e..2f2fea0424 100644
--- a/crates/ruff/src/rules/pyupgrade/rules/use_pep604_annotation.rs
+++ b/crates/ruff/src/rules/pyupgrade/rules/use_pep604_annotation.rs
@@ -60,6 +60,7 @@ pub(crate) fn use_pep604_annotation(
     // Avoid fixing forward references, or types not in an annotation.
     let fixable = checker.semantic().in_type_definition()
         && !checker.semantic().in_complex_string_type_definition();
+
     match operator {
         Pep604Operator::Optional => {
             let mut diagnostic = Diagnostic::new(NonPEP604Annotation, expr.range());
diff --git a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP007.py.snap b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP007.py.snap
index dc112cd0e0..250930ff96 100644
--- a/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP007.py.snap
+++ b/crates/ruff/src/rules/pyupgrade/snapshots/ruff__rules__pyupgrade__tests__UP007.py.snap
@@ -276,4 +276,20 @@ UP007.py:60:8: UP007 [*] Use `X | Y` for type annotations
    60 |+    x: str | int
 61 61 |     x: Union["str", "int"]
 
+UP007.py:61:8: UP007 [*] Use `X | Y` for type annotations
+   |
+59 |     x = Union["str", "int"]
+60 |     x: Union[str, int]
+61 |     x: Union["str", "int"]
+   |        ^^^^^^^^^^^^^^^^^^^ UP007
+   |
+   = help: Convert to `X | Y`
+
+ℹ Suggested fix
+58 58 |     x = Union[str, int]
+59 59 |     x = Union["str", "int"]
+60 60 |     x: Union[str, int]
+61    |-    x: Union["str", "int"]
+   61 |+    x: "str" | "int"
+
 
diff --git a/crates/ruff_python_semantic/src/analyze/typing.rs b/crates/ruff_python_semantic/src/analyze/typing.rs
index 4da230cd70..39ea75af86 100644
--- a/crates/ruff_python_semantic/src/analyze/typing.rs
+++ b/crates/ruff_python_semantic/src/analyze/typing.rs
@@ -127,22 +127,36 @@ pub fn to_pep604_operator(
     slice: &Expr,
     semantic: &SemanticModel,
 ) -> Option<Pep604Operator> {
-    /// Returns `true` if any argument in the slice is a string.
-    fn any_arg_is_str(slice: &Expr) -> bool {
+    /// Returns `true` if any argument in the slice is a quoted annotation).
+    fn quoted_annotation(slice: &Expr) -> bool {
         match slice {
             Expr::Constant(ast::ExprConstant {
                 value: Constant::Str(_),
                 ..
             }) => true,
-            Expr::Tuple(ast::ExprTuple { elts, .. }) => elts.iter().any(any_arg_is_str),
+            Expr::Tuple(ast::ExprTuple { elts, .. }) => elts.iter().any(quoted_annotation),
             _ => false,
         }
     }
 
-    // If any of the _arguments_ are forward references, we can't use PEP 604.
-    // Ex) `Union["str", "int"]` can't be converted to `"str" | "int"`.
-    if any_arg_is_str(slice) {
-        return None;
+    // If the slice is a forward reference (e.g., `Optional["Foo"]`), it can only be rewritten
+    // if we're in a typing-only context.
+    //
+    // This, for example, is invalid, as Python will evaluate `"Foo" | None` at runtime in order to
+    // populate the function's `__annotations__`:
+    // ```python
+    // def f(x: "Foo" | None): ...
+    // ```
+    //
+    // This, however, is valid:
+    // ```python
+    // def f():
+    //     x: "Foo" | None
+    // ```
+    if quoted_annotation(slice) {
+        if semantic.execution_context().is_runtime() {
+            return None;
+        }
     }
 
     semantic