openrewrite-migrate-python 0.14.0.dev20260901135710__tar.gz → 0.14.0.dev20260901153140__tar.gz

{openrewrite_migrate_python-0.14.0.dev20260901135710 → openrewrite_migrate_python-0.14.0.dev20260901153140}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openrewrite-migrate-python
-Version: 0.14.0.dev20260901135710
+Version: 0.14.0.dev20260901153140
 Summary: OpenRewrite recipes for migrating Python codebases to newer Python versions.
 Author-email: "Moderne Inc." <support@moderne.io>
 License: Moderne Proprietary

{openrewrite_migrate_python-0.14.0.dev20260901135710 → openrewrite_migrate_python-0.14.0.dev20260901153140}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "openrewrite-migrate-python"
 description = "OpenRewrite recipes for migrating Python codebases to newer Python versions."
-version = "0.14.0.dev20260901135710"  # Updated dynamically during release
+version = "0.14.0.dev20260901153140"  # Updated dynamically during release
 authors = [{ name = "Moderne Inc.", email = "support@moderne.io" }]
 license = { text = "Moderne Proprietary" }
 readme = "README.md"

openrewrite_migrate_python-0.14.0.dev20260901153140/src/openrewrite_migrate_python/migrate/_typing_imports.py

@@ -0,0 +1,164 @@
+"""Internal helpers for locating references to the typing module in a file."""
+
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple, cast
+
+from rewrite.java.support_types import JRightPadded, Space, Statement
+from rewrite.java.tree import Block, If, Import, J
+from rewrite.python.import_utils import get_alias_name, get_name_string, get_qualid_name
+from rewrite.python.remove_import import RemoveImport, RemoveImportOptions, prefix_to_inherit
+from rewrite.python.tree import CompilationUnit, MultiImport
+from rewrite.python.visitor import PythonVisitor
+
+
+def _module_scope_statements(statements: Sequence[Statement]) -> Iterator[Statement]:
+    """Every statement binding names in the module scope, `if` bodies included.
+
+    `if TYPE_CHECKING:` is where files that defer their annotations keep their
+    typing imports.
+    """
+    for stmt in statements:
+        yield stmt
+        body = _unconditional_body(stmt) if isinstance(stmt, If) else None
+        if body is not None:
+            yield from _module_scope_statements(body.statements)
+
+
+def _unconditional_body(if_: If) -> Optional[Block]:
+    """The body of an `if` that only adds bindings to the enclosing scope.
+
+    None once there is an `else`: the branches are then alternative bindings of
+    the same name, and honouring one would rewrite the other's binding too.
+    """
+    then_part = if_.then_part
+    return then_part if if_.else_part is None and isinstance(then_part, Block) else None
+
+
+def scan_typing_imports(cu: CompilationUnit) -> Tuple[Dict[str, str], Set[str]]:
+    """Map how this file can reference typing members.
+
+    Returns:
+        A tuple of:
+        - exposed name -> typing member, for ``from typing import member [as alias]``
+        - the set of names bound to the typing module itself, for
+          ``import typing [as alias]``
+    """
+    members: Dict[str, str] = {}
+    module_refs: Set[str] = set()
+    for stmt in _module_scope_statements(cu.statements):
+        if isinstance(stmt, MultiImport):
+            if stmt.from_ is not None:
+                if get_name_string(stmt.from_) == "typing":
+                    for imp in stmt.names:
+                        name = get_qualid_name(imp.qualid)
+                        members.setdefault(get_alias_name(imp) or name, name)
+            else:
+                for imp in stmt.names:
+                    if get_qualid_name(imp.qualid) == "typing":
+                        module_refs.add(get_alias_name(imp) or "typing")
+        elif isinstance(stmt, Import):
+            if get_qualid_name(stmt.qualid) == "typing":
+                module_refs.add(get_alias_name(stmt) or "typing")
+    return members, module_refs
+
+
+def maybe_remove_typing_import(visitor: PythonVisitor, name: Optional[str] = None) -> None:
+    """Schedule removal of a typing import once it is no longer used.
+
+    Mirrors the SDK's `maybe_remove_import`, registering the block-aware visitor.
+    """
+    options = RemoveImportOptions(module="typing", name=name)
+    if visitor._after_visit is None:
+        visitor._after_visit = []
+    for scheduled in visitor._after_visit:
+        # Only this visitor counts as already scheduled: a plain `RemoveImport`
+        # for the same name reaches the top level alone, and would leave the
+        # nested imports this exists to remove.
+        if (
+            type(scheduled) is _RemoveImportInBlocks
+            and scheduled.module == options.module
+            and scheduled.name == options.name
+        ):
+            return
+    visitor._after_visit.append(_RemoveImportInBlocks(options))
+
+
+def maybe_remove_typing_module_import(visitor: PythonVisitor, module_refs: Set[str]) -> None:
+    """Schedule removal of `import typing` once it is no longer used.
+
+    RemoveImport's usage check counts `typing` identifiers, which an aliased
+    module import doesn't produce — removal is only scheduled when an
+    unaliased `import typing` exists.
+    """
+    if "typing" in module_refs:
+        maybe_remove_typing_import(visitor)
+
+
+class _RemoveImportInBlocks(RemoveImport):
+    """`RemoveImport` extended to imports inside module-scope `if` bodies.
+
+    The base visitor prunes only the file's top-level statement list.
+    """
+
+    def visit_compilation_unit(self, cu: CompilationUnit, p: Any) -> J:
+        cu = cast(CompilationUnit, super().visit_compilation_unit(cu, p))
+        kept = self._prune(cu.padding.statements, prune_imports=False)
+        return cu if kept is None else cu.padding.replace(_statements=kept)
+
+    def _prune(
+        self, padded_statements: Sequence[JRightPadded], prune_imports: bool
+    ) -> Optional[List[JRightPadded]]:
+        """The statements with the import gone, or None when nothing changed."""
+        kept: List[JRightPadded] = []
+        inherited: Optional[Space] = None
+        changed = False
+        for index, padded in enumerate(padded_statements):
+            stmt = padded.element
+            if isinstance(stmt, If):
+                result: Optional[Statement] = self._prune_if(stmt)
+            elif prune_imports and isinstance(stmt, MultiImport):
+                result = self._process_multi_import(stmt)
+            elif prune_imports and isinstance(stmt, Import):
+                result = self._process_single_import(stmt)
+            else:
+                result = stmt
+
+            if result is None:
+                changed = True
+                if inherited is None:
+                    inherited = prefix_to_inherit(stmt, index)
+                continue
+            if result is not stmt:
+                padded = JRightPadded(result, padded.after, padded.markers)
+                changed = True
+            if inherited is not None:
+                if not padded.element.prefix.comments:
+                    padded = JRightPadded(
+                        padded.element.replace(prefix=inherited), padded.after, padded.markers
+                    )
+                elif inherited.comments:
+                    # Neither comment can be given up, and one prefix holds both.
+                    return None
+                inherited = None
+            kept.append(padded)
+        # A leftover prefix means nothing followed the removal to carry it, so a
+        # comment on the removed statement would be dropped with it.
+        if inherited is not None and inherited.comments:
+            return None
+        return kept if changed else None
+
+    def _prune_if(self, if_: If) -> Optional[If]:
+        """`if_` with the import gone from its body, or None once that empties it."""
+        then_part = _unconditional_body(if_)
+        if then_part is None:
+            return if_
+        kept = self._prune(then_part.padding.statements, prune_imports=True)
+        if kept is None:
+            return if_
+        if not kept:
+            return None
+        padded = if_.padding.then_part
+        return if_.padding.replace(
+            _then_part=JRightPadded(
+                then_part.padding.replace(_statements=kept), padded.after, padded.markers
+            )
+        )

{openrewrite_migrate_python-0.14.0.dev20260901135710 → openrewrite_migrate_python-0.14.0.dev20260901153140}/src/openrewrite_migrate_python/migrate/typing_deprecations.py

@@ -20,18 +20,21 @@ handled with `ChangeImport` in `upgrade_to_python39`.
 See: https://peps.python.org/pep-0585/
 """
 
-import re
+import ast
 from typing import Any, ClassVar, Optional, Set, cast
 
 from rewrite import ExecutionContext, Recipe, TreeVisitor
 from rewrite.category import CategoryDescriptor
 from rewrite.decorators import categorize
 from rewrite.marketplace import Python
-from rewrite.python.import_utils import get_alias_name, get_name_string, get_qualid_name
-from rewrite.python.remove_import import RemoveImportOptions, maybe_remove_import
+from rewrite.python.markers import Quoted
 from rewrite.python.scope_utils import LocalBindings
 
-from ._typing_imports import maybe_remove_typing_module_import
+from ._typing_imports import (
+    maybe_remove_typing_import,
+    maybe_remove_typing_module_import,
+    scan_typing_imports,
+)
 from rewrite.python.tree import CompilationUnit, MultiImport, NamedArgument
 from rewrite.python.visitor import PythonVisitor
 from rewrite.java.tree import (
@@ -78,67 +81,78 @@ def _make_builtin_alias_editor(old_name: str, new_name: str) -> "TreeVisitor[Any
             self._bare_ref: Optional[str] = None
             # Names that refer to the typing module ('typing' or import aliases).
             self._module_refs: Set[str] = set()
-            self._keep_from_import: bool = False
             self._rewrote_qualified: bool = False
             self._local_bindings = LocalBindings()
 
         def visit_compilation_unit(
             self, cu: CompilationUnit, p: ExecutionContext
         ) -> CompilationUnit:
-            self._bare_ref = None
-            self._module_refs = set()
             self._rewrote_qualified = False
             self._local_bindings = LocalBindings()
 
-            for stmt in cu.statements:
-                if isinstance(stmt, MultiImport):
-                    if stmt.from_ is not None:
-                        if get_name_string(stmt.from_) == "typing":
-                            for imp in stmt.names:
-                                if get_qualid_name(imp.qualid) == old_name and self._bare_ref is None:
-                                    self._bare_ref = get_alias_name(imp) or old_name
-                    else:
-                        for imp in stmt.names:
-                            if get_qualid_name(imp.qualid) == "typing":
-                                self._module_refs.add(get_alias_name(imp) or "typing")
-                elif isinstance(stmt, Import):
-                    if get_qualid_name(stmt.qualid) == "typing":
-                        self._module_refs.add(get_alias_name(stmt) or "typing")
+            members, self._module_refs = scan_typing_imports(cu)
+            self._bare_ref = next(
+                (exposed for exposed, name in members.items() if name == old_name), None
+            )
 
             if self._bare_ref is None and not self._module_refs:
                 return cu
 
-            # A reference inside a string annotation is a Literal, not an
-            # Identifier, so it is neither renamed nor seen by RemoveImport's
-            # usage check; keep the import so such references stay resolvable.
-            self._keep_from_import = self._bare_ref is not None and self._used_in_string_literal(
-                cu, self._bare_ref
-            )
-
             result = cast(CompilationUnit, super().visit_compilation_unit(cu, p))
 
-            if self._bare_ref is not None and not self._keep_from_import:
-                maybe_remove_import(
-                    self, RemoveImportOptions(module="typing", name=old_name)
-                )
+            # A string forward reference is neither renamed above nor seen by
+            # RemoveImport's usage check; scanning the rewritten tree finds the
+            # ones that outlived the rename.
+            if self._bare_ref is not None and not self._forward_reference_survives(
+                result, self._bare_ref
+            ):
+                maybe_remove_typing_import(self, old_name)
             if self._rewrote_qualified:
                 maybe_remove_typing_module_import(self, self._module_refs)
 
             return result
 
         @staticmethod
-        def _used_in_string_literal(cu: CompilationUnit, name: str) -> bool:
-            pattern = re.compile(rf"\b{re.escape(name)}\b")
+        def _forward_reference_survives(cu: CompilationUnit, name: str) -> bool:
+            """True when a string forward reference still names ``name``.
+
+            The parser leaves such a reference as text -- a string `Literal`
+            annotation, or an `Identifier` marked `Quoted` -- so its contents are
+            parsed as an expression. Prose that merely spells the name out, as a
+            docstring does, is not an expression and so is not a reference.
+            """
             found = False
 
-            class LiteralScan(PythonVisitor[Any]):
+            def names_target(text: str) -> bool:
+                if name not in text:
+                    return False
+                try:
+                    parsed = ast.parse(text.strip(), mode="eval")
+                except (SyntaxError, ValueError):
+                    return False
+                return any(
+                    isinstance(node, ast.Name) and node.id == name
+                    for node in ast.walk(parsed)
+                )
+
+            class ForwardReferenceScan(PythonVisitor[Any]):
                 def visit_literal(self, literal: Literal, p: Any) -> J:
                     nonlocal found
-                    if not found and isinstance(literal.value, str) and pattern.search(literal.value):
+                    if not found and isinstance(literal.value, str) and names_target(literal.value):
                         found = True
                     return literal
 
-            LiteralScan().visit(cu, None)
+                def visit_identifier(self, ident: Identifier, p: Any) -> J:
+                    nonlocal found
+                    if (
+                        not found
+                        and ident.markers.find_first(Quoted) is not None
+                        and names_target(ident.simple_name)
+                    ):
+                        found = True
+                    return ident
+
+            ForwardReferenceScan().visit(cu, None)
             return found
 
         def visit_field_access(self, field_access: FieldAccess, p: ExecutionContext) -> J:
@@ -203,7 +217,10 @@ class _ReplaceTypingAliasRecipe(Recipe):
     def description(self) -> str:
         return (
             f"PEP 585 deprecated `typing.{self.old_name}` in Python 3.9. "
-            f"Replace with the built-in `{self.new_name}` type for generic annotations."
+            f"Replace with the built-in `{self.new_name}` type for generic annotations. "
+            "Applies inside `if TYPE_CHECKING:` blocks as well; the built-in generic is "
+            "valid at runtime from Python 3.9 on, so no `from __future__ import "
+            "annotations` is needed."
         )
 
     def editor(self) -> TreeVisitor[Any, ExecutionContext]:

{openrewrite_migrate_python-0.14.0.dev20260901135710 → openrewrite_migrate_python-0.14.0.dev20260901153140}/src/openrewrite_migrate_python/migrate/typing_union_to_pipe.py

@@ -18,7 +18,6 @@ from rewrite.decorators import categorize
 from rewrite.marketplace import Python
 from rewrite.markers import Markers
 from rewrite.python.markers import Quoted
-from rewrite.python.remove_import import RemoveImportOptions, maybe_remove_import
 from rewrite.python.visitor import PythonVisitor
 from rewrite.java.support_types import JRightPadded, Space
 from rewrite.java.tree import (
@@ -33,7 +32,11 @@ from rewrite.java.tree import (
 from rewrite.python.tree import CollectionLiteral, CompilationUnit, UnionType
 from rewrite.utils import random_id
 
-from ._typing_imports import maybe_remove_typing_module_import, scan_typing_imports
+from ._typing_imports import (
+    maybe_remove_typing_import,
+    maybe_remove_typing_module_import,
+    scan_typing_imports,
+)
 
 # Define category path: Python > Migrate > Python 3.10
 _Python310 = [
@@ -89,7 +92,7 @@ class _PipeSyntaxVisitor(PythonVisitor[ExecutionContext]):
         if isinstance(clazz, FieldAccess):
             maybe_remove_typing_module_import(self, self._module_refs)
         else:
-            maybe_remove_import(self, RemoveImportOptions(module="typing", name=member))
+            maybe_remove_typing_import(self, member)
 
 
 def _create_none_literal(prefix: Space) -> Literal:
@@ -171,7 +174,10 @@ class ReplaceTypingOptionalWithUnion(Recipe):
     def description(self) -> str:
         return (
             "PEP 604 introduced the `|` operator for union types in Python 3.10. "
-            "Replace `Optional[X]` with the more concise `X | None` syntax."
+            "Replace `Optional[X]` with the more concise `X | None` syntax. "
+            "Applies inside `if TYPE_CHECKING:` blocks as well; `X | None` is valid at "
+            "runtime from Python 3.10 on, so no `from __future__ import annotations` "
+            "is needed."
         )
 
     @property
@@ -258,7 +264,10 @@ class ReplaceTypingUnionWithPipe(Recipe):
     def description(self) -> str:
         return (
             "PEP 604 introduced the `|` operator for union types in Python 3.10. "
-            "Replace `Union[X, Y, ...]` with the more concise `X | Y | ...` syntax."
+            "Replace `Union[X, Y, ...]` with the more concise `X | Y | ...` syntax. "
+            "Applies inside `if TYPE_CHECKING:` blocks as well; `X | Y` is valid at "
+            "runtime from Python 3.10 on, so no `from __future__ import annotations` "
+            "is needed."
         )
 
     @property

{openrewrite_migrate_python-0.14.0.dev20260901135710 → openrewrite_migrate_python-0.14.0.dev20260901153140}/src/openrewrite_migrate_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openrewrite-migrate-python
-Version: 0.14.0.dev20260901135710
+Version: 0.14.0.dev20260901153140
 Summary: OpenRewrite recipes for migrating Python codebases to newer Python versions.
 Author-email: "Moderne Inc." <support@moderne.io>
 License: Moderne Proprietary

{openrewrite_migrate_python-0.14.0.dev20260901135710 → openrewrite_migrate_python-0.14.0.dev20260901153140}/tests/test_typing_deprecations.py

@@ -430,6 +430,66 @@ class TestPep585CornerCases:
             )
         )
 
+    def test_removes_import_when_only_a_docstring_mentions_the_name(self):
+        """Prose that happens to spell the name is not a reference to it."""
+        spec = RecipeSpec(recipe=ReplaceTypingListWithList())
+        spec.rewrite_run(
+            python(
+                """
+                from typing import List
+
+
+                def f(items):
+                    'Return a List of things.'
+                    return items
+
+
+                x: List[int] = []
+                """,
+                """
+                def f(items):
+                    'Return a List of things.'
+                    return items
+
+
+                x: list[int] = []
+                """,
+            )
+        )
+
+    def test_keeps_import_used_by_quoted_annotation(self):
+        """A quoted annotation the rename leaves alone still needs the import."""
+        spec = RecipeSpec(recipe=ReplaceTypingListWithList())
+        spec.rewrite_run(
+            python(
+                """
+                from typing import List
+                x: List[int] = []
+                y: "List[int]" = []
+                """,
+                """
+                from typing import List
+                x: list[int] = []
+                y: "List[int]" = []
+                """,
+            )
+        )
+
+    def test_removes_import_when_the_quoted_annotation_is_renamed(self):
+        """The rename reaches inside a bare quoted annotation, freeing the import."""
+        spec = RecipeSpec(recipe=ReplaceTypingListWithList())
+        spec.rewrite_run(
+            python(
+                """
+                from typing import List
+                x: "List" = []
+                """,
+                """
+                x: "list" = []
+                """,
+            )
+        )
+
     def test_does_not_rename_keyword_argument(self):
         """A keyword argument named `List` binds a parameter, not the import."""
         spec = RecipeSpec(recipe=ReplaceTypingListWithList())
@@ -521,3 +581,194 @@ class TestReplaceTypingText:
                 """
             )
         )
+
+
+class TestTypeCheckingBlocks:
+    """Typing imports nested in an `if TYPE_CHECKING:` block."""
+
+    def test_rewrites_and_drops_emptied_block(self):
+        spec = RecipeSpec(recipe=ReplaceTypingListWithList())
+        spec.rewrite_run(
+            python(
+                """
+                from __future__ import annotations
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    from typing import List
+
+                def f(x: List[int]) -> None: ...
+                """,
+                """
+                from __future__ import annotations
+                from typing import TYPE_CHECKING
+
+                def f(x: list[int]) -> None: ...
+                """,
+            )
+        )
+
+    def test_keeps_block_holding_other_imports(self):
+        spec = RecipeSpec(recipe=ReplaceTypingListWithList())
+        spec.rewrite_run(
+            python(
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    from typing import List, Optional
+
+                def f(x: List[int], y: Optional[str]) -> None: ...
+                """,
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    from typing import Optional
+
+                def f(x: list[int], y: Optional[str]) -> None: ...
+                """,
+            )
+        )
+
+    def test_rewrites_qualified_reference(self):
+        spec = RecipeSpec(recipe=ReplaceTypingListWithList())
+        spec.rewrite_run(
+            python(
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    import typing
+
+                def f(x: typing.List[int]) -> None: ...
+                """,
+                """
+                from typing import TYPE_CHECKING
+
+                def f(x: list[int]) -> None: ...
+                """,
+            )
+        )
+
+    def test_ignores_a_block_whose_else_binds_the_name(self):
+        """Renaming the `else` binding too would rebind the builtin."""
+        spec = RecipeSpec(recipe=ReplaceTypingListWithList())
+        spec.rewrite_run(
+            python(
+                """
+                from typing import TYPE_CHECKING
+                import collections
+
+                if TYPE_CHECKING:
+                    from typing import List
+                else:
+                    List = collections.UserList
+
+                def f(x: List[int]) -> None: ...
+                """
+            )
+        )
+
+    def test_keeps_a_comment_on_the_statement_after_a_removed_import(self):
+        spec = RecipeSpec(recipe=ReplaceTypingListWithList())
+        spec.rewrite_run(
+            python(
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    from typing import List
+                    # needed for annotations
+                    from typing import Dict
+
+                def f(x: List[int], y: Dict[str, int]) -> None: ...
+                """,
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    # needed for annotations
+                    from typing import Dict
+
+                def f(x: list[int], y: Dict[str, int]) -> None: ...
+                """,
+            )
+        )
+
+    def test_keeps_import_whose_comment_would_be_lost(self):
+        """No prefix can hold the comment once the import it leads is gone."""
+        RecipeSpec(recipe=ReplaceTypingListWithList()).rewrite_run(
+            python(
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    # only needed for annotations
+                    from typing import List
+
+                def f(x: List[int]) -> None: ...
+                """,
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    # only needed for annotations
+                    from typing import List
+
+                def f(x: list[int]) -> None: ...
+                """,
+            )
+        )
+
+        RecipeSpec(recipe=ReplaceTypingListWithList()).rewrite_run(
+            python(
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    from collections.abc import Sequence
+                    # only needed for annotations
+                    from typing import List
+
+                def f(x: List[int], y: Sequence[int]) -> None: ...
+                """,
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    from collections.abc import Sequence
+                    # only needed for annotations
+                    from typing import List
+
+                def f(x: list[int], y: Sequence[int]) -> None: ...
+                """,
+            )
+        )
+
+        RecipeSpec(recipe=ReplaceTypingListWithList()).rewrite_run(
+            python(
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    # comment A
+                    from typing import List
+                    # comment B
+                    from typing import Dict
+
+                def f(x: List[int], y: Dict[str, int]) -> None: ...
+                """,
+                """
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    # comment A
+                    from typing import List
+                    # comment B
+                    from typing import Dict
+
+                def f(x: list[int], y: Dict[str, int]) -> None: ...
+                """,
+            )
+        )

{openrewrite_migrate_python-0.14.0.dev20260901135710 → openrewrite_migrate_python-0.14.0.dev20260901153140}/tests/test_typing_union_to_pipe.py

@@ -340,3 +340,29 @@ class TestReplaceTypingUnionWithPipe:
                 """,
             )
         )
+
+
+class TestTypeCheckingBlocks:
+    """Typing imports nested in an `if TYPE_CHECKING:` block."""
+
+    def test_rewrites_optional_and_drops_emptied_block(self):
+        spec = RecipeSpec(recipe=ReplaceTypingOptionalWithUnion())
+        spec.rewrite_run(
+            python(
+                """
+                from __future__ import annotations
+                from typing import TYPE_CHECKING
+
+                if TYPE_CHECKING:
+                    from typing import Optional
+
+                def f(y: Optional[str]) -> None: ...
+                """,
+                """
+                from __future__ import annotations
+                from typing import TYPE_CHECKING
+
+                def f(y: str | None) -> None: ...
+                """,
+            )
+        )