Avoid raising PEP 604 errors with forward-referenced members (#3640)

astral-sh · Mar 21, 2023 · 626169e · 626169e
1 parent e9f359a
commit 626169e
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 28 deletions.
diff --git a/crates/ruff/resources/test/fixtures/pyupgrade/UP007.py b/crates/ruff/resources/test/fixtures/pyupgrade/UP007.py
@@ -46,3 +46,8 @@ def f(x: Union[("str", "int"), float]) -> None:
 def f() -> None:
     x: Optional[str]
     x = Optional[str]
+
+    x = Union[str, int]
+    x = Union["str", "int"]
+    x: Union[str, int]
+    x: Union["str", "int"]
diff --git a/crates/ruff/src/rules/pyupgrade/rules/use_pep604_annotation.rs b/crates/ruff/src/rules/pyupgrade/rules/use_pep604_annotation.rs
@@ -80,6 +80,12 @@ enum TypingMember {
 
 /// UP007
 pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, slice: &Expr) {
+    // 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;
+    }
+
     let Some(typing_member) = checker.ctx.resolve_call_path(value).as_ref().and_then(|call_path| {
         if checker.ctx.match_typing_call_path(call_path, "Optional") {
             Some(TypingMember::Optional)
@@ -92,8 +98,8 @@ pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, s
         return;
     };
 
-    // Avoid fixing forward annotations.
-    let fixable = !checker.ctx.in_deferred_string_type_definition && !any_arg_is_str(slice);
+    // Avoid fixing forward references.
+    let fixable = !checker.ctx.in_deferred_string_type_definition;
 
     match typing_member {
         TypingMember::Optional => {

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
@@ -181,32 +181,6 @@ expression: diagnostics
     column: 33
   fix: ~
   parent: ~
-- kind:
-    name: NonPEP604Annotation
-    body: "Use `X | Y` for type annotations"
-    suggestion: ~
-    fixable: false
-  location:
-    row: 38
-    column: 9
-  end_location:
-    row: 38
-    column: 26
-  fix: ~
-  parent: ~
-- kind:
-    name: NonPEP604Annotation
-    body: "Use `X | Y` for type annotations"
-    suggestion: ~
-    fixable: false
-  location:
-    row: 42
-    column: 9
-  end_location:
-    row: 42
-    column: 37
-  fix: ~
-  parent: ~
 - kind:
     name: NonPEP604Annotation
     body: "Use `X | Y` for type annotations"
@@ -227,4 +201,24 @@ expression: diagnostics
       row: 47
       column: 20
   parent: ~
+- kind:
+    name: NonPEP604Annotation
+    body: "Use `X | Y` for type annotations"
+    suggestion: "Convert to `X | Y`"
+    fixable: true
+  location:
+    row: 52
+    column: 7
+  end_location:
+    row: 52
+    column: 22
+  fix:
+    content: str | int
+    location:
+      row: 52
+      column: 7
+    end_location:
+      row: 52
+      column: 22
+  parent: ~