Python: imports in an if TYPE_CHECKING: block are module scope

[knutwannheden]

ChangeImport and RemoveImport decided what a file imports by iterating cu.statements. An import nested in if TYPE_CHECKING: is a child of a J.If at that level, not a statement of the compilation unit, so neither visitor ever saw it.

org.openrewrite.python.migrate.ReplaceTypingCallableWithCollectionsAbcCallable is ChangeImport(old_module="typing", old_name="Callable", new_module="collections.abc") and nothing else. It reported no change at all on this file:

from __future__ import annotations
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import Callable

def f(x: Callable[[int], str]) -> None: ...

The typing-alias moves in UpgradeToPython39 are the same recipe under different arguments and skipped the same files. A file that defers its annotations keeps its typing imports exactly here, so this is not a corner of the input space.

Where the replacement import goes

ChangeImport binds the new import inside the block it took the old one from, rather than letting maybe_add_import place a replacement at module level. Hoisting it would turn an import the file deliberately deferred into one that runs at import time.

A nested match is spliced where it stands: the reduced old statement keeps its place and the new import follows it, or merges into a sibling from <new_module> import … already in the block. That merge is what keeps a chain of alias moves from emitting one line per alias. Four UpgradeToPython39 moves over

if TYPE_CHECKING:
    from typing import Deque, Iterable, Pattern, Sequence

leave

if TYPE_CHECKING:
    from re import Pattern
    from collections.abc import Iterable, Sequence
    from collections import deque

A block can hold more than one match — from typing import Callable and import typing together — and each gets its binding.

What counts as module scope

unconditional_body is the rule: an if carrying an else is not module scope. Its branches are alternative bindings of the same name, and honouring one rewrites the binding the other was there to provide.

Everything below that boundary is left exactly as it stands, and visit_import and visit_multi_import both require the statement's parent to be the compilation unit:

• a function or class body, where pruning would strip a binding and put nothing back
• a try:, while: or with: body, or an else/elif branch, where pruning empties the suite and the file stops parsing

References follow the same line. visit_identifier renames wherever it goes, and a suite is not a scope, so an else-branch import reads as a module global that LocalBindings will not shadow. When the recipe changes the bound name and such a match survives, the file is left whole rather than renamed onto a name nothing binds.

Comments

RemoveImport abandons a removal rather than lose a comment, at three points:

• nothing follows the removed statement to carry one
• the statement that would inherit its prefix has one of its own
• a second removal has one the first prefix cannot also hold

ChangeImport skips the merge when the replaced statement carries a comment, so the comment stays with the import it describes.

Java side

PythonRemoveImportVisitor.mightChange transcribes the Python visitor's early returns so the host can skip a round trip. Left alone it would answer "no" for precisely the files this fixes, and the RPC path would keep the old behaviour. It excludes an if carrying an else, matching module_scope_blocks.

Tests

Across tests/python/test_remove_import.py and tests/recipes/test_change_import.py, both arms of the existing native/java fixture. Each guard is mutation-pinned: reverting it fails the test named for it.

Verification

rewrite-migrate-python installed against this branch: its suite passes on main (602) and on the open PR #89 (612), and the reported no-op above now rewrites in place.

Note for whoever bumps that repo: RemoveImport._prune is now _prune_statements. PR #89's _RemoveImportInBlocks — a caller-side subclass that existed only because the SDK could not do this — overrides the old name, and should be deleted in the same bump along with its maybe_remove_typing_import wrapper.