diff --git a/rewrite-python/rewrite/src/rewrite/python/add_import.py b/rewrite-python/rewrite/src/rewrite/python/add_import.py index 4a07266946..ee87635060 100644 --- a/rewrite-python/rewrite/src/rewrite/python/add_import.py +++ b/rewrite-python/rewrite/src/rewrite/python/add_import.py @@ -253,8 +253,8 @@ def _try_merge_into_existing(self, cu: CompilationUnit) -> CompilationUnit: # case-insensitive alphabetical position (mirrors rewrite-javascript's # AddImport). Existing members are not reordered; only the new member # is positioned. - new_import = self._create_import_element(self.name, self.alias) - existing_padded = self._insert_member( + new_import = create_import_element(self.name, self.alias) + existing_padded = insert_member( list(stmt.padding.names.padding.elements), new_import) # Recreate the MultiImport with the new names @@ -279,66 +279,9 @@ def _try_merge_into_existing(self, cu: CompilationUnit) -> CompilationUnit: return cu - def _insert_member(self, elements: list, new_import: Import) -> list: - """Insert ``new_import`` into a list of ``JRightPadded[Import]`` at its - case-insensitive alphabetical position. - - The space after the ``import`` keyword lives in the surrounding - ``JContainer.before``, so the element at index 0 carries an empty prefix - while every later element carries a single-space prefix (the space after - the separating comma). Trailing whitespace (e.g. before a ``)`` in a - parenthesized import) lives in the last element's ``.after`` and must - travel with whichever element ends up last. - """ - insert_idx = self._sorted_insert_index(elements, new_import) - end = len(elements) - - if insert_idx == 0: - prefix = Space.EMPTY - if elements: - # The displaced first element now follows a comma. - first = elements[0] - elements[0] = first.replace( - _element=first.element.replace(prefix=Space.SINGLE_SPACE)) - else: - prefix = Space.SINGLE_SPACE - new_import = new_import.replace(prefix=prefix) - - if insert_idx == end and elements: - # Appending at the end: the new element becomes the last, so any - # trailing whitespace moves from the old last element onto it. - last = elements[-1] - elements[-1] = last.replace(_after=Space.EMPTY) - after = last.after - else: - after = Space.EMPTY - - elements.insert(insert_idx, JRightPadded(new_import, after, Markers.EMPTY)) - return elements - - def _sorted_insert_index(self, elements: list, new_import: Import) -> int: - """Return the index at which the new member keeps the list in - case-insensitive alphabetical order: the first existing member whose - bound name sorts after the new member's, or the end if none does. - - Members are sorted by their bound name (alias if present, else the - imported name), matching rewrite-javascript's comparator. - """ - new_key = self._sort_key(new_import) - for i, padded in enumerate(elements): - if new_key < self._sort_key(padded.element): - return i - return len(elements) - - @staticmethod - def _sort_key(imp: Import) -> str: - """Case-insensitive sort key for an imported member: its alias if it has - one, otherwise the imported name.""" - return (get_alias_name(imp) or get_qualid_name(imp.qualid)).lower() - def _add_import(self, cu: CompilationUnit) -> CompilationUnit: """Add a new import statement to the compilation unit.""" - new_import = self._create_multi_import() + new_import = create_import_statement(self.module, self.name, self.alias) # Insert after the module docstring (which must stay first) and any existing imports. padded_stmts = list(cu.padding.statements) @@ -383,133 +326,191 @@ def _add_import(self, cu: CompilationUnit) -> CompilationUnit: return cu.padding.replace(_statements=padded_stmts) - def _create_multi_import(self) -> MultiImport: - """Create a new MultiImport statement.""" - if self.name is None: - # Direct import: import module [as alias] - import_elem = self._create_import_element(self.module, self.alias) - return MultiImport( - random_id(), - Space([], '\n'), - Markers.EMPTY, - None, # No 'from' - False, # Not parenthesized - JContainer( - Space.SINGLE_SPACE, - [pad_right(import_elem)], - Markers.EMPTY - ) - ) - else: - # From import: from module import name [as alias] - from_name = self._create_module_name(self.module) - # Add space prefix (the space between 'from' and module name) - from_name = from_name.replace(prefix=Space.SINGLE_SPACE) - import_elem = self._create_import_element(self.name, self.alias) - return MultiImport( - random_id(), - Space([], '\n'), - Markers.EMPTY, - JRightPadded(from_name, Space.SINGLE_SPACE, Markers.EMPTY), - False, # Not parenthesized - JContainer( - Space.SINGLE_SPACE, - [pad_right(import_elem)], - Markers.EMPTY - ) - ) - def _create_import_element(self, name: str, alias: Optional[str]) -> Import: - """Create an Import element.""" - qualid = self._create_qualified_name(name) - alias_left_padded = None - if alias: - alias_ident = Identifier( - random_id(), +def insert_member(elements: list, new_import: Import) -> list: + """Insert ``new_import`` into a list of ``JRightPadded[Import]`` at its + case-insensitive alphabetical position. + + The space after the ``import`` keyword lives in the surrounding + ``JContainer.before``, so the element at index 0 carries an empty prefix + while every later element carries a single-space prefix (the space after + the separating comma). Trailing whitespace (e.g. before a ``)`` in a + parenthesized import) lives in the last element's ``.after`` and must + travel with whichever element ends up last. + """ + insert_idx = _sorted_insert_index(elements, new_import) + end = len(elements) + + if insert_idx == 0: + prefix = Space.EMPTY + if elements: + # The displaced first element now follows a comma. + first = elements[0] + elements[0] = first.replace( + _element=first.element.replace(prefix=Space.SINGLE_SPACE)) + else: + prefix = Space.SINGLE_SPACE + new_import = new_import.replace(prefix=prefix) + + if insert_idx == end and elements: + # Appending at the end: the new element becomes the last, so any + # trailing whitespace moves from the old last element onto it. + last = elements[-1] + elements[-1] = last.replace(_after=Space.EMPTY) + after = last.after + else: + after = Space.EMPTY + + elements.insert(insert_idx, JRightPadded(new_import, after, Markers.EMPTY)) + return elements + +def _sorted_insert_index(elements: list, new_import: Import) -> int: + """Return the index at which the new member keeps the list in + case-insensitive alphabetical order: the first existing member whose + bound name sorts after the new member's, or the end if none does. + + Members are sorted by their bound name (alias if present, else the + imported name), matching rewrite-javascript's comparator. + """ + new_key = _sort_key(new_import) + for i, padded in enumerate(elements): + if new_key < _sort_key(padded.element): + return i + return len(elements) + +def _sort_key(imp: Import) -> str: + """Case-insensitive sort key for an imported member: its alias if it has + one, otherwise the imported name.""" + return (get_alias_name(imp) or get_qualid_name(imp.qualid)).lower() + + +def create_import_statement(module: str, name: Optional[str] = None, + alias: Optional[str] = None) -> MultiImport: + """Build ``import module [as alias]`` or ``from module import name [as alias]``.""" + if name is None: + # Direct import: import module [as alias] + import_elem = create_import_element(module, alias) + return MultiImport( + random_id(), + Space([], '\n'), + Markers.EMPTY, + None, # No 'from' + False, # Not parenthesized + JContainer( Space.SINGLE_SPACE, - Markers.EMPTY, - [], - alias, - None, - None + [pad_right(import_elem)], + Markers.EMPTY ) - alias_left_padded = JLeftPadded( + ) + else: + # From import: from module import name [as alias] + from_name = _create_module_name(module) + # Add space prefix (the space between 'from' and module name) + from_name = from_name.replace(prefix=Space.SINGLE_SPACE) + import_elem = create_import_element(name, alias) + return MultiImport( + random_id(), + Space([], '\n'), + Markers.EMPTY, + JRightPadded(from_name, Space.SINGLE_SPACE, Markers.EMPTY), + False, # Not parenthesized + JContainer( Space.SINGLE_SPACE, - alias_ident, + [pad_right(import_elem)], Markers.EMPTY ) + ) - return Import( +def create_import_element(name: str, alias: Optional[str] = None) -> Import: + """Build one member of an import statement: ``name [as alias]``.""" + qualid = _create_qualified_name(name) + alias_left_padded = None + if alias: + alias_ident = Identifier( random_id(), - Space.EMPTY, + Space.SINGLE_SPACE, Markers.EMPTY, - JLeftPadded(Space.EMPTY, False, Markers.EMPTY), - qualid, - alias_left_padded + [], + alias, + None, + None + ) + alias_left_padded = JLeftPadded( + Space.SINGLE_SPACE, + alias_ident, + Markers.EMPTY ) - def _create_qualified_name(self, name: str) -> FieldAccess: - """Create a FieldAccess for a qualified name.""" - parts = name.split('.') - if len(parts) == 1: - return FieldAccess( - random_id(), + return Import( + random_id(), + Space.EMPTY, + Markers.EMPTY, + JLeftPadded(Space.EMPTY, False, Markers.EMPTY), + qualid, + alias_left_padded + ) + +def _create_qualified_name(name: str) -> FieldAccess: + """Create a FieldAccess for a qualified name.""" + parts = name.split('.') + if len(parts) == 1: + return FieldAccess( + random_id(), + Space.EMPTY, + Markers.EMPTY, + Empty(random_id(), Space.EMPTY, Markers.EMPTY), + JLeftPadded( Space.EMPTY, - Markers.EMPTY, - Empty(random_id(), Space.EMPTY, Markers.EMPTY), - JLeftPadded( - Space.EMPTY, - Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None), - Markers.EMPTY - ), - None - ) + Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None), + Markers.EMPTY + ), + None + ) - # Build nested FieldAccess for qualified names like "os.path" - # Start with the first part as an Identifier target - result: J = Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None) - # Wrap remaining parts as FieldAccess nodes - for part in parts[1:]: - result = FieldAccess( - random_id(), + # Build nested FieldAccess for qualified names like "os.path" + # Start with the first part as an Identifier target + result: J = Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None) + # Wrap remaining parts as FieldAccess nodes + for part in parts[1:]: + result = FieldAccess( + random_id(), + Space.EMPTY, + Markers.EMPTY, + result, + JLeftPadded( Space.EMPTY, - Markers.EMPTY, - result, - JLeftPadded( - Space.EMPTY, - Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], part, None, None), - Markers.EMPTY - ), - None - ) - # For multi-part names, result is already a FieldAccess. - # Wrap single-part Identifier in FieldAccess(Empty, name) for consistency - # (but single-part is handled above, so this shouldn't happen) - assert isinstance(result, FieldAccess) - return result - - def _create_module_name(self, name: str) -> J: - """Create a name tree for use as a 'from' module name. - - Single-part names become Identifier; multi-part become FieldAccess. - Unlike _create_qualified_name, this does NOT wrap single parts in - FieldAccess(Empty, name) since the 'from' printer has no special - handling for Empty targets. - """ - parts = name.split('.') - result: J = Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None) - for part in parts[1:]: - result = FieldAccess( - random_id(), + Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], part, None, None), + Markers.EMPTY + ), + None + ) + # For multi-part names, result is already a FieldAccess. + # Wrap single-part Identifier in FieldAccess(Empty, name) for consistency + # (but single-part is handled above, so this shouldn't happen) + assert isinstance(result, FieldAccess) + return result + +def _create_module_name(name: str) -> J: + """Create a name tree for use as a 'from' module name. + + Single-part names become Identifier; multi-part become FieldAccess. + Unlike _create_qualified_name, this does NOT wrap single parts in + FieldAccess(Empty, name) since the 'from' printer has no special + handling for Empty targets. + """ + parts = name.split('.') + result: J = Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None) + for part in parts[1:]: + result = FieldAccess( + random_id(), + Space.EMPTY, + Markers.EMPTY, + result, + JLeftPadded( Space.EMPTY, - Markers.EMPTY, - result, - JLeftPadded( - Space.EMPTY, - Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], part, None, None), - Markers.EMPTY - ), - None - ) - return result - + Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], part, None, None), + Markers.EMPTY + ), + None + ) + return result diff --git a/rewrite-python/rewrite/src/rewrite/python/import_utils.py b/rewrite-python/rewrite/src/rewrite/python/import_utils.py index 580cac49e8..b5d161568f 100644 --- a/rewrite-python/rewrite/src/rewrite/python/import_utils.py +++ b/rewrite-python/rewrite/src/rewrite/python/import_utils.py @@ -15,14 +15,37 @@ """Shared utility functions for Python import handling.""" import ast -from typing import Optional, Tuple +from typing import Iterator, Optional, Sequence, Tuple -from rewrite.java.support_types import JavaType, JRightPadded, Space -from rewrite.java.tree import Empty, FieldAccess, Identifier, Import +from rewrite.java.support_types import JavaType, JRightPadded, Space, Statement +from rewrite.java.tree import Block, Empty, FieldAccess, Identifier, If, Import from rewrite.markers import Markers from rewrite.python.markers import CanonicalName, Quoted +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 module_scope_blocks(statements: Sequence[Statement]) -> Iterator[Block]: + """The `if` bodies whose bindings land in the module scope. + + `if TYPE_CHECKING:` is where files that defer their annotations keep their + typing imports. + """ + for stmt in statements: + body = unconditional_body(stmt) if isinstance(stmt, If) else None + if body is not None: + yield body + yield from module_scope_blocks(body.statements) + + def get_qualid_name(qualid) -> str: """Get the string representation of a qualified name.""" if isinstance(qualid, Identifier): diff --git a/rewrite-python/rewrite/src/rewrite/python/recipes/change_import.py b/rewrite-python/rewrite/src/rewrite/python/recipes/change_import.py index 61dd1be536..c98daa07b6 100644 --- a/rewrite-python/rewrite/src/rewrite/python/recipes/change_import.py +++ b/rewrite-python/rewrite/src/rewrite/python/recipes/change_import.py @@ -15,7 +15,7 @@ """Recipe to change Python imports from one module/name to another.""" from dataclasses import dataclass, field, replace as dc_replace -from typing import Any, Optional +from typing import Any, List, Optional, Tuple from rewrite import ExecutionContext, Recipe, TreeVisitor from rewrite.category import CategoryDescriptor @@ -23,18 +23,24 @@ from rewrite.marketplace import Python from rewrite.recipe import option from rewrite.java import J -from rewrite.java.support_types import JavaType -from rewrite.java.tree import FieldAccess, Identifier, Import, MethodInvocation -from rewrite.python.import_utils import get_qualid_name, get_name_string, get_alias_name +from rewrite.java.support_types import JavaType, JContainer, JRightPadded, Statement +from rewrite.java.tree import FieldAccess, Identifier, If, Import, MethodInvocation, Space +from rewrite.markers import Markers +from rewrite.python.import_utils import (get_qualid_name, get_name_string, get_alias_name, + module_scope_blocks, unconditional_body) from rewrite.python.scope_utils import LocalBindings from rewrite.python.tree import CompilationUnit, MultiImport from rewrite.python.visitor import PythonVisitor -from rewrite.python.add_import import AddImportOptions, maybe_add_import +from rewrite.python.add_import import (AddImportOptions, create_import_element, + create_import_statement, insert_member, maybe_add_import) from rewrite.python.remove_import import RemoveImportOptions, maybe_remove_import, prefix_to_inherit _Imports = [*Python, CategoryDescriptor(display_name="Imports")] +# An import to bind, as (module, name, alias); a None name means `import module`. +_Binding = Tuple[str, Optional[str], Optional[str]] + def _create_module_type(fqn: str) -> JavaType.Class: """Create a JavaType.Class for a module from its fully qualified name. @@ -144,6 +150,8 @@ class ChangeImportVisitor(PythonVisitor[ExecutionContext]): rewrote_qualified_refs: bool = False new_module_type: Optional[JavaType.Class] = None local_bindings: LocalBindings # a fresh instance per compilation unit + old_import_at_module_level: bool = False + direct_module_import_at_module_level: bool = False def visit_compilation_unit(self, cu: CompilationUnit, p: ExecutionContext) -> J: self.has_old_import = False @@ -154,35 +162,20 @@ def visit_compilation_unit(self, cu: CompilationUnit, p: ExecutionContext) -> J: self.new_module_type = None self.local_bindings = LocalBindings() - # Single pass: detect old imports and direct module imports for stmt in cu.statements: - if isinstance(stmt, Import) and not isinstance(stmt, MultiImport): - if not self.has_old_import: - alias = self._check_for_old_single_import(stmt) - if alias is not None: - self.has_old_import = True - self.old_alias = alias if alias != "" else None - if old_name and not self.has_direct_module_import: - name = get_qualid_name(stmt.qualid) - if name == old_module: - self.has_direct_module_import = True - self.module_alias = get_alias_name(stmt) - elif isinstance(stmt, MultiImport): - if not self.has_old_import: - alias = self._check_for_old_import(stmt) - if alias is not None: - self.has_old_import = True - self.old_alias = alias if alias != "" else None - if old_name and not self.has_direct_module_import and stmt.from_ is None: - for imp in stmt.names: - name = get_qualid_name(imp.qualid) - if name == old_module: - self.has_direct_module_import = True - self.module_alias = get_alias_name(imp) - break + self._detect(stmt) + # Where the old import is found decides where the replacement goes. + self.old_import_at_module_level = self.has_old_import + self.direct_module_import_at_module_level = self.has_direct_module_import + for block in module_scope_blocks(cu.statements): + for stmt in block.statements: + self._detect(stmt) if not self.has_old_import and not self.has_direct_module_import: return cu + if old_name and (new_alias or self.old_alias or new_name) != \ + (self.old_alias or old_name) and self._match_outside_module_scope(cu): + return cu # Visit to transform imports result = super().visit_compilation_unit(cu, p) @@ -190,9 +183,9 @@ def visit_compilation_unit(self, cu: CompilationUnit, p: ExecutionContext) -> J: return result result = self._transfer_removed_prefixes(cu, result) + result = self._rewrite_block_imports(result) - # Schedule adding the new import (only for direct import changes) - if self.has_old_import: + if self.old_import_at_module_level: alias_to_use = new_alias or self.old_alias if new_name: maybe_add_import(self, AddImportOptions( @@ -210,21 +203,74 @@ def visit_compilation_unit(self, cu: CompilationUnit, p: ExecutionContext) -> J: # If we rewrote qualified references, manage the direct import if self.rewrote_qualified_refs: - maybe_add_import(self, AddImportOptions( - module=new_module, - alias=new_alias, - only_if_referenced=False - )) + if self.direct_module_import_at_module_level: + maybe_add_import(self, AddImportOptions( + module=new_module, + alias=new_alias, + only_if_referenced=False + )) maybe_remove_import(self, RemoveImportOptions( module=old_module, )) return result + def _detect(self, stmt: Statement) -> None: + """Record what `stmt` binds: the import being changed, and the module whose + qualified references would be rewritten.""" + if isinstance(stmt, MultiImport): + if not self.has_old_import: + alias = self._check_for_old_import(stmt) + if alias is not None: + self.has_old_import = True + self.old_alias = alias if alias != "" else None + if old_name and not self.has_direct_module_import and stmt.from_ is None: + for imp in stmt.names: + if get_qualid_name(imp.qualid) == old_module: + self.has_direct_module_import = True + self.module_alias = get_alias_name(imp) + break + elif isinstance(stmt, Import): + if not self.has_old_import: + alias = self._check_for_old_single_import(stmt) + if alias is not None: + self.has_old_import = True + self.old_alias = alias if alias != "" else None + if old_name and not self.has_direct_module_import: + if get_qualid_name(stmt.qualid) == old_module: + self.has_direct_module_import = True + self.module_alias = get_alias_name(stmt) + + def _match_outside_module_scope(self, cu: CompilationUnit) -> bool: + """True when a match sits somewhere this recipe leaves alone. That import + goes on binding the old name, so renaming the references it serves would + leave them unresolved.""" + in_scope = {stmt.id for stmt in cu.statements} + for block in module_scope_blocks(cu.statements): + in_scope.update(stmt.id for stmt in block.statements) + found: List[bool] = [] + outer = self + + class Finder(PythonVisitor): + def visit_multi_import(self, multi: MultiImport, p) -> J: + if (multi.id not in in_scope and + outer._check_for_old_import(multi) is not None): + found.append(True) + return multi + + Finder().visit(cu, None) + return bool(found) + + def _at_module_level(self) -> bool: + """True for a statement of the compilation unit. Replacements are bound here + or, by `_rewrite_block`, in a module-scope `if` body; a match deeper than that + would be removed with nothing put in its place.""" + return isinstance(self.cursor.parent_tree_cursor().value, CompilationUnit) + def visit_import(self, import_: Import, p: ExecutionContext) -> Optional[J]: # ty: ignore[invalid-method-override] if not self.has_old_import or old_name: return import_ - if self.cursor.first_enclosing(MultiImport): + if not self._at_module_level(): return import_ alias = self._check_for_old_single_import(import_) if alias is None: @@ -234,6 +280,8 @@ def visit_import(self, import_: Import, p: ExecutionContext) -> Optional[J]: # def visit_multi_import(self, multi: MultiImport, p: ExecutionContext) -> Optional[J]: # ty: ignore[invalid-method-override] if not self.has_old_import: return multi + if not self._at_module_level(): + return multi alias = self._check_for_old_import(multi) if alias is None: @@ -353,6 +401,143 @@ def visit_field_access(self, field_access: FieldAccess, p: ExecutionContext) -> result = result.padding.replace(_name=result.padding.name.replace(_element=new_name_ident)) return result + def _rewrite_block_imports(self, cu: CompilationUnit) -> CompilationUnit: + """Rewrite a match inside an `if TYPE_CHECKING:`-style block where it stands. + + The replacement import is bound in the same block. Hoisting it to module level + — what maybe_add_import would do — would run at import time an import the file + deliberately deferred. + """ + kept = self._rewrite_statements(cu.padding.statements) + return cu if kept is None else cu.padding.replace(_statements=kept) + + def _rewrite_statements(self, padded_statements) -> Optional[List[JRightPadded]]: + """`padded_statements` with every module-scope `if` body rewritten, or None + when none of them held a match.""" + kept: List[JRightPadded] = [] + changed = False + for padded in padded_statements: + stmt = padded.element + if isinstance(stmt, If): + rewritten = self._rewrite_if(stmt) + if rewritten is not stmt: + padded = padded.replace(_element=rewritten) + changed = True + kept.append(padded) + return kept if changed else None + + def _rewrite_if(self, if_: If) -> If: + body = unconditional_body(if_) + if body is None: + return if_ + kept = self._rewrite_block(body.padding.statements) + if kept is None: + return if_ + padded = if_.padding.then_part + return if_.padding.replace(_then_part=JRightPadded( + body.padding.replace(_statements=kept), padded.after, padded.markers)) + + def _rewrite_block(self, padded_statements) -> Optional[List[JRightPadded]]: + """The block's statements with every match replaced by the new import, or None + when nothing in it matched. Nested `if` bodies are rewritten too.""" + kept: List[JRightPadded] = [] + changed = False + to_add: List[Tuple[_Binding, int, Space]] = [] + for padded in padded_statements: + stmt = padded.element + if isinstance(stmt, If): + rewritten = self._rewrite_if(stmt) + if rewritten is not stmt: + padded = padded.replace(_element=rewritten) + changed = True + kept.append(padded) + continue + reduced, binding = self._match_in_block(stmt) + if reduced is not stmt: + changed = True + if reduced is not None: + kept.append(padded if reduced is stmt else padded.replace(_element=reduced)) + if binding is not None: + # A comment on the statement describes what it imports, so it + # travels only when the whole statement is replaced. + prefix = (stmt.prefix if reduced is None + else Space([], stmt.prefix.whitespace)) + to_add.append((binding, len(kept), prefix)) + # Back to front, so an insertion never shifts a pending position. + for binding, at, prefix in reversed(to_add): + if self._place_import(kept, binding, at, prefix): + changed = True + return kept if changed else None + + def _match_in_block(self, stmt: Statement) -> Tuple[Optional[Statement], + Optional[_Binding]]: + """`(statement to keep, import to bind here)` for a match in a block, and + `(stmt, None)` for anything else.""" + if isinstance(stmt, MultiImport): + alias = self._check_for_old_import(stmt) + if alias is not None: + bound = new_alias or (alias or None) + if old_name: + return (self._remove_name_from_import(stmt, old_name), + (new_module, new_name, bound)) + return (self._remove_module_from_import(stmt, old_module), + (new_module, None, bound)) + elif isinstance(stmt, Import): + alias = self._check_for_old_single_import(stmt) + if alias is not None: + # The module is the whole statement, so the statement goes. + return None, (new_module, None, new_alias or (alias or None)) + # `import old_module` behind references this recipe rewrote to the new module: + # bind the new module here as well, and let RemoveImport drop the old one once + # nothing refers to it. + if (old_name and self.rewrote_qualified_refs and + not self.direct_module_import_at_module_level and + self._binds_module(stmt, old_module)): + return stmt, (new_module, None, new_alias) + return stmt, None + + @staticmethod + def _binds_module(stmt: Statement, module: str) -> bool: + """True when `stmt` is an `import module` rather than a `from` import.""" + if isinstance(stmt, MultiImport): + return stmt.from_ is None and any( + get_qualid_name(imp.qualid) == module for imp in stmt.names) + return isinstance(stmt, Import) and get_qualid_name(stmt.qualid) == module + + def _place_import(self, kept: List[JRightPadded], binding: _Binding, + at: int, prefix: Space) -> bool: + """Bind `binding` in the block, merged into a sibling import from the same + module when there is one and otherwise as a statement of its own at `at`. + False when the block already binds it.""" + module, name, alias = binding + if name is None: + if any(self._binds_module(p.element, module) for p in kept): + return False + else: + bound = alias or name + for index, padded in enumerate(kept): + stmt = padded.element + if not isinstance(stmt, MultiImport) or stmt.from_ is None: + continue + if get_name_string(stmt.from_) != module: + continue + elements = list(stmt.padding.names.padding.elements) + if any((get_alias_name(e.element) or get_qualid_name(e.element.qualid)) + == bound for e in elements): + return False + if prefix.comments: + # A merge has nowhere to carry that comment. + break + kept[index] = padded.replace(_element=stmt.padding.replace( + _names=JContainer(stmt.padding.names.before, + insert_member(elements, + create_import_element(name, alias)), + stmt.padding.names.markers))) + return True + statement = create_import_statement(module, name, alias).replace(prefix=prefix) + kept.insert(at, JRightPadded(statement, Space.EMPTY, Markers.EMPTY)) + return True + def _transfer_removed_prefixes(self, before: CompilationUnit, after: CompilationUnit) -> CompilationUnit: """Dropping a statement discards its prefix; when it is worth rescuing (see prefix_to_inherit) hand it to the next surviving statement, @@ -429,7 +614,8 @@ def _check_for_old_import(self, multi: MultiImport) -> Optional[str]: return get_alias_name(imp) or "" return None - def _remove_name_from_import(self, multi: MultiImport, name_to_remove: str) -> Optional[J]: + def _remove_name_from_import(self, multi: MultiImport, + name_to_remove: str) -> Optional[MultiImport]: """Remove a specific name from a 'from X import a, b, c' statement.""" from rewrite.java.support_types import JContainer from rewrite.java.tree import Space @@ -456,7 +642,8 @@ def _remove_name_from_import(self, multi: MultiImport, name_to_remove: str) -> O ) return multi - def _remove_module_from_import(self, multi: MultiImport, module_to_remove: str) -> Optional[J]: + def _remove_module_from_import(self, multi: MultiImport, + module_to_remove: str) -> Optional[MultiImport]: """Remove a module from an import statement.""" from rewrite.java.support_types import JContainer from rewrite.java.tree import Space diff --git a/rewrite-python/rewrite/src/rewrite/python/remove_import.py b/rewrite-python/rewrite/src/rewrite/python/remove_import.py index 992837193c..8e989e67bd 100644 --- a/rewrite-python/rewrite/src/rewrite/python/remove_import.py +++ b/rewrite-python/rewrite/src/rewrite/python/remove_import.py @@ -15,13 +15,14 @@ """RemoveImport visitor for Python import handling.""" from dataclasses import dataclass -from typing import Optional, Set +from typing import List, Optional, Sequence, Set from rewrite.java import J -from rewrite.java.support_types import JContainer, JRightPadded -from rewrite.java.tree import Identifier, Import, Space +from rewrite.java.support_types import JContainer, JRightPadded, Statement +from rewrite.java.tree import Identifier, If, Import, Space from rewrite.python.import_utils import (get_qualid_name, get_name_string, get_alias_name, - get_canonical_fqn, referenced_names) + get_canonical_fqn, referenced_names, + unconditional_body) from rewrite.python.scope_utils import LocalBindings from rewrite.python.tree import CompilationUnit, MultiImport from rewrite.python.visitor import PythonVisitor @@ -179,59 +180,68 @@ def visit_identifier(self, ident: Identifier, p) -> J: def _remove_import(self, cu: CompilationUnit) -> CompilationUnit: """Remove the import from the compilation unit.""" - new_padded_stmts = [] + kept = self._prune_statements(cu.padding.statements) + return cu if kept is None else cu.padding.replace(_statements=kept) + + def _prune_statements(self, padded_statements: Sequence[JRightPadded]) -> Optional[List[JRightPadded]]: + """The statements with the import gone, or None to leave them as they are — + because nothing matched, or because a comment would go with the removal. One + prefix cannot carry two comments, so a collision abandons the removal.""" + kept: List[JRightPadded] = [] + inherited: Optional[Space] = None changed = False - removed_prefix = None - - for index, padded in enumerate(cu.padding.statements): + for index, padded in enumerate(padded_statements): stmt = padded.element - if isinstance(stmt, Import) and not isinstance(stmt, MultiImport): - result = self._process_single_import(stmt) - if result is None: - removed_prefix = prefix_to_inherit(stmt, index) - changed = True - else: - if removed_prefix is not None: - padded = JRightPadded( - padded.element.replace(prefix=removed_prefix), - padded.after, padded.markers - ) - removed_prefix = None - new_padded_stmts.append(padded) - continue if isinstance(stmt, MultiImport): - result = self._process_multi_import(stmt) - if result is None: - # Remove the entire statement; remember its prefix - # so the next statement can inherit it if needed - removed_prefix = prefix_to_inherit(stmt, index) - changed = True - else: - if result is not stmt: - padded = JRightPadded(result, padded.after, padded.markers) - changed = True - # Transfer removed statement's prefix to this statement - if removed_prefix is not None: - padded = JRightPadded( - padded.element.replace(prefix=removed_prefix), - padded.after, padded.markers - ) - removed_prefix = None - new_padded_stmts.append(padded) + result: Optional[Statement] = self._process_multi_import(stmt) + elif isinstance(stmt, Import): + result = self._process_single_import(stmt) + elif isinstance(stmt, If): + result = self._prune_if_body(stmt) else: - # Transfer removed statement's prefix to the next statement - if removed_prefix is not None: - new_padded_stmts.append(JRightPadded( - stmt.replace(prefix=removed_prefix), - padded.after, padded.markers - )) - removed_prefix = None - else: - new_padded_stmts.append(padded) - - if changed: - return cu.padding.replace(_statements=new_padded_stmts) - return cu + result = stmt + + if result is None: + changed = True + if inherited is None: + inherited = prefix_to_inherit(stmt, index) + elif stmt.prefix.comments: + return None + 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: + 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_body(self, if_: If) -> Optional[If]: + """`if_` with the import gone from its body, or None once that empties it.""" + body = unconditional_body(if_) + if body is None: + return if_ + kept = self._prune_statements(body.padding.statements) + if kept is None: + return if_ + if not kept: + return None + padded = if_.padding.then_part + return if_.padding.replace( + _then_part=JRightPadded( + body.padding.replace(_statements=kept), padded.after, padded.markers + ) + ) def _process_single_import(self, imp: Import) -> Optional[Import]: """Process a standalone J.Import. Return None to remove, or the original.""" diff --git a/rewrite-python/rewrite/tests/python/test_remove_import.py b/rewrite-python/rewrite/tests/python/test_remove_import.py index 3d2814c50f..ce53886355 100644 --- a/rewrite-python/rewrite/tests/python/test_remove_import.py +++ b/rewrite-python/rewrite/tests/python/test_remove_import.py @@ -623,6 +623,156 @@ def f(x=clock): ) +class TestImportsInBlocks: + """Imports nested in a module-scope `if` body, where `if TYPE_CHECKING:` keeps them.""" + + def test_remove_nested_import_and_drop_emptied_block(self, arm): + spec = RecipeSpec(recipe=from_visitor( + _remove_import_visitor(arm, 'typing', 'List', only_if_unused=False))) + spec.rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import List + + x = 1 + """, + """ + from typing import TYPE_CHECKING + + x = 1 + """, + ) + ) + + def test_keep_block_holding_other_imports(self, arm): + spec = RecipeSpec(recipe=from_visitor( + _remove_import_visitor(arm, 'typing', 'List', only_if_unused=False))) + spec.rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import List + from os.path import join + + x = 1 + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from os.path import join + + x = 1 + """, + ) + ) + + def test_keep_import_when_emptying_would_lose_a_comment(self, arm): + spec = RecipeSpec(recipe=from_visitor( + _remove_import_visitor(arm, 'typing', 'List', only_if_unused=False))) + spec.rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + # only needed for annotations + from typing import List + + x = 1 + """ + ) + ) + + def test_keep_imports_when_a_later_removal_would_lose_a_comment(self, arm): + spec = RecipeSpec(recipe=from_visitor( + _remove_import_visitor(arm, 'typing', 'List', only_if_unused=False))) + spec.rewrite_run( + python( + """ + # header + from typing import List + # about the second one + from typing import List + + x = 1 + """ + ) + ) + + def test_keep_imports_when_the_next_statement_has_its_own_comment(self, arm): + spec = RecipeSpec(recipe=from_visitor( + _remove_import_visitor(arm, 'typing', 'List', only_if_unused=False))) + spec.rewrite_run( + python( + """ + # about the List import + from typing import List + # about x + x = 1 + """ + ) + ) + + def test_keep_import_when_the_block_has_an_else(self, arm): + spec = RecipeSpec(recipe=from_visitor( + _remove_import_visitor(arm, 'typing', 'List', only_if_unused=False))) + spec.rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import List + else: + List = list + + x = 1 + """ + ) + ) + + def test_keep_import_shadowed_only_by_a_type_checking_binding(self, arm): + """The block binding does not exist at run time, so it shadows nothing.""" + spec = RecipeSpec(recipe=from_visitor( + _remove_import_visitor(arm, 'typing', 'List'))) + spec.rewrite_run( + python( + """ + from typing import TYPE_CHECKING + from typing import List + + if TYPE_CHECKING: + from mymod import List + + def f() -> List[int]: + return List() + """ + ) + ) + + def test_keep_nested_import_that_is_still_used(self, arm): + spec = RecipeSpec(recipe=from_visitor( + _remove_import_visitor(arm, 'typing', 'List'))) + spec.rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import List + + def f(x: List[int]) -> None: ... + """ + ) + ) + + class TestRemoveImportStringAnnotations: """A string annotation is a forward reference, so the names inside it are load-bearing: whatever resolves the annotation later needs their imports.""" diff --git a/rewrite-python/rewrite/tests/recipes/test_change_import.py b/rewrite-python/rewrite/tests/recipes/test_change_import.py index 8066b059f8..a8ff6c358c 100644 --- a/rewrite-python/rewrite/tests/recipes/test_change_import.py +++ b/rewrite-python/rewrite/tests/recipes/test_change_import.py @@ -1027,3 +1027,340 @@ def f(): perf_counter = 1 """, ) + + +class TestImportsInBlocks: + """Imports nested in a module-scope `if` body, where `if TYPE_CHECKING:` keeps them.""" + + @staticmethod + def _callable_to_abc() -> RecipeSpec: + return RecipeSpec(recipe=ChangeImport( + old_module='typing', + old_name='Callable', + new_module='collections.abc', + )) + + def test_sole_member_is_rewritten_in_place(self): + self._callable_to_abc().rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import Callable + + def f(x: Callable[[int], str]) -> None: ... + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from collections.abc import Callable + + def f(x: Callable[[int], str]) -> None: ... + """, + ) + ) + + RecipeSpec(recipe=ChangeImport( + old_module='collections', + old_name='Mapping', + new_module='collections.abc', + new_name='Map', + )).rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from collections import Mapping + + def f(x: Mapping) -> None: ... + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from collections.abc import Map + + def f(x: Map) -> None: ... + """, + ) + ) + + def test_split_leaves_the_new_import_in_the_block(self): + self._callable_to_abc().rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import Callable, Optional + + def f(x: Callable[[int], str], y: Optional[str]) -> None: ... + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import Optional + from collections.abc import Callable + + def f(x: Callable[[int], str], y: Optional[str]) -> None: ... + """, + ) + ) + + def test_merges_into_an_existing_import_in_the_same_block(self): + """What keeps a sequence of alias moves from emitting one line per alias.""" + self._callable_to_abc().rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import Callable, Optional + from collections.abc import Sequence + + def f(x: Callable[[int], str], y: Optional[Sequence[str]]) -> None: ... + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import Optional + from collections.abc import Callable, Sequence + + def f(x: Callable[[int], str], y: Optional[Sequence[str]]) -> None: ... + """, + ) + ) + + def test_nested_direct_import_is_rewritten_in_place(self): + spec = RecipeSpec(recipe=ChangeImport( + old_module='urllib2', + new_module='urllib.request', + )) + spec.rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + import urllib2 + + x = 1 + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + import urllib.request + + x = 1 + """, + ) + ) + + def test_a_match_outside_module_scope_is_left_alone(self): + """Each file needs a module-scope match too, or the recipe returns before it looks.""" + self._callable_to_abc().rewrite_run( + python( + """ + import sys + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import Callable + + if sys.version_info >= (3, 10): + pass + else: + from typing import Callable + + def f(x: Callable[[int], str]) -> None: ... + """, + """ + import sys + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from collections.abc import Callable + + if sys.version_info >= (3, 10): + pass + else: + from typing import Callable + + def f(x: Callable[[int], str]) -> None: ... + """, + ), + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import Callable + + def g(): + from typing import Callable + return Callable + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from collections.abc import Callable + + def g(): + from typing import Callable + return Callable + """, + ), + ) + + def test_qualified_reference_binds_the_new_module_in_the_block(self): + self._callable_to_abc().rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + import typing + + def f(x: typing.Callable[[int], str]) -> None: ... + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + import collections.abc + + def f(x: collections.abc.Callable[[int], str]) -> None: ... + """, + ) + ) + + def test_qualified_reference_does_not_duplicate_an_existing_import(self): + self._callable_to_abc().rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + import typing + import collections.abc + + def f(x: typing.Callable[[int], str], y: collections.abc.Sequence) -> None: ... + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + import collections.abc + + def f(x: collections.abc.Callable[[int], str], y: collections.abc.Sequence) -> None: ... + """, + ) + ) + + def test_a_rename_stops_when_a_match_survives_out_of_scope(self): + RecipeSpec(recipe=ChangeImport( + old_module='collections', + old_name='Mapping', + new_module='collections.abc', + new_name='Map', + )).rewrite_run( + python( + """ + import sys + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from collections import Mapping + + if sys.version_info >= (3, 10): + pass + else: + from collections import Mapping + m = Mapping() + """ + ) + ) + + def test_both_import_forms_in_one_block_are_rewritten(self): + self._callable_to_abc().rewrite_run( + python( + """ + from __future__ import annotations + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from typing import Callable + import typing + + def f(x: Callable[[int], str], y: typing.Callable[[int], int]) -> None: ... + """, + """ + from __future__ import annotations + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from collections.abc import Callable + import collections.abc + + def f(x: Callable[[int], str], y: collections.abc.Callable[[int], int]) -> None: ... + """, + ) + ) + + def test_a_comment_on_the_replaced_statement_keeps_it_out_of_a_merge(self): + self._callable_to_abc().rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from collections.abc import Sequence + # only needed for annotations + from typing import Callable + + def f(x: Callable[[int], str], y: Sequence) -> None: ... + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from collections.abc import Sequence + # only needed for annotations + from collections.abc import Callable + + def f(x: Callable[[int], str], y: Sequence) -> None: ... + """, + ) + ) + + def test_split_leaves_the_comment_on_the_statement_it_describes(self): + self._callable_to_abc().rewrite_run( + python( + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + # only needed for annotations + from typing import Callable, Optional + + def f(x: Callable[[int], str], y: Optional[str]) -> None: ... + """, + """ + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + # only needed for annotations + from typing import Optional + from collections.abc import Callable + + def f(x: Callable[[int], str], y: Optional[str]) -> None: ... + """, + ) + ) diff --git a/rewrite-python/src/main/java/org/openrewrite/python/service/PythonRemoveImportVisitor.java b/rewrite-python/src/main/java/org/openrewrite/python/service/PythonRemoveImportVisitor.java index cff87fe6df..5c7f3abcd8 100644 --- a/rewrite-python/src/main/java/org/openrewrite/python/service/PythonRemoveImportVisitor.java +++ b/rewrite-python/src/main/java/org/openrewrite/python/service/PythonRemoveImportVisitor.java @@ -25,7 +25,9 @@ import org.openrewrite.java.tree.Statement; import org.openrewrite.python.tree.Py; +import java.util.ArrayList; import java.util.HashMap; +import java.util.List; import java.util.Map; import static org.openrewrite.python.internal.PythonImportNames.canonicalFqn; @@ -69,7 +71,7 @@ Map options() { @Override boolean mightChange(Py.CompilationUnit cu) { - for (Statement statement : cu.getStatements()) { + for (Statement statement : moduleScopeStatements(cu.getStatements())) { if (statement instanceof J.Import) { if (name == null && module.equals(nameString(((J.Import) statement).getQualid()))) { return true; @@ -81,6 +83,23 @@ boolean mightChange(Py.CompilationUnit cu) { return false; } + /** + * Every statement binding names in the module scope, {@code if} bodies included. Mirrors + * {@code module_scope_blocks} on the Python side, including its exclusion of an {@code if} + * that carries an {@code else}. + */ + private static List moduleScopeStatements(List statements) { + List flattened = new ArrayList<>(statements); + for (Statement statement : statements) { + if (statement instanceof J.If && ((J.If) statement).getElsePart() == null && + ((J.If) statement).getThenPart() instanceof J.Block) { + flattened.addAll(moduleScopeStatements( + ((J.Block) ((J.If) statement).getThenPart()).getStatements())); + } + } + return flattened; + } + /** * Whether any member of the statement is one the Python implementation would consider removing. * {@code only_if_unused} is not evaluated here: it can only retain a candidate, so leaving it