fix Optional["ForwardRef"] rewriting

asottile · Nov 10, 2022 · 9796546 · 9796546
1 parent 5c27928
commit 9796546
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 10 deletions.
diff --git a/pyupgrade/_plugins/typing_pep604.py b/pyupgrade/_plugins/typing_pep604.py
@@ -144,16 +144,17 @@ def visit_Subscript(
     if not _supported_version(state):
         return
 
-    # prevent rewriting forward annotations
-    if (
-            (sys.version_info >= (3, 9) and _any_arg_is_str(node.slice)) or
-            (
-                sys.version_info < (3, 9) and
-                isinstance(node.slice, ast.Index) and
-                _any_arg_is_str(node.slice.value)
-            )
-    ):
-        return
+    # don't rewrite forward annotations (unless we know they will be dequoted)
+    if 'annotations' not in state.from_imports['__future__']:
+        if (
+                (sys.version_info >= (3, 9) and _any_arg_is_str(node.slice)) or
+                (
+                    sys.version_info < (3, 9) and
+                    isinstance(node.slice, ast.Index) and
+                    _any_arg_is_str(node.slice.value)
+                )
+        ):
+            return
 
     if is_name_attr(
             node.value,

diff --git a/tests/features/typing_pep604_test.py b/tests/features/typing_pep604_test.py
@@ -185,6 +185,17 @@ def f(x: int | str) -> None: ...
 
             id='Optional rewrite multi-line',
         ),
+        pytest.param(
+            'from __future__ import annotations\n'
+            'from typing import Optional\n'
+            'x: Optional["str"]\n',
+
+            'from __future__ import annotations\n'
+            'from typing import Optional\n'
+            'x: str | None\n',
+
+            id='Optional rewrite with forward reference',
+        ),
         pytest.param(
             'from typing import Union, Sequence\n'
             'def f(x: Union[Union[A, B], Sequence[Union[C, D]]]): pass\n',