diff --git a/libcst/codemod/commands/rename.py b/libcst/codemod/commands/rename.py index aad4cea6..9ad4334d 100644 --- a/libcst/codemod/commands/rename.py +++ b/libcst/codemod/commands/rename.py @@ -103,15 +103,20 @@ def as_name(self, value: Optional[Tuple[str, str]]) -> None: self.context.scratch["as_name"] = value @property - def scheduled_removals(self) -> Set[cst.CSTNode]: + def scheduled_removals( + self, + ) -> Set[Union[cst.CSTNode, Tuple[str, Optional[str], Optional[str]]]]: """A set of nodes that have been renamed to help with the cleanup of now potentially unused - imports, during import cleanup in `leave_Module`.""" + imports, during import cleanup in `leave_Module`. Can also contain tuples that can be passed + directly to RemoveImportsVisitor.remove_unused_import().""" if "scheduled_removals" not in self.context.scratch: self.context.scratch["scheduled_removals"] = set() return self.context.scratch["scheduled_removals"] @scheduled_removals.setter - def scheduled_removals(self, value: Set[cst.CSTNode]) -> None: + def scheduled_removals( + self, value: Set[Union[cst.CSTNode, Tuple[str, Optional[str], Optional[str]]]] + ) -> None: self.context.scratch["scheduled_removals"] = value @property @@ -169,11 +174,13 @@ def leave_Import( and import_alias.asname is not None ): self.bypass_import = True - # TODO: put this into self.scheduled_removals - RemoveImportsVisitor.remove_unused_import( - self.context, - import_alias.evaluated_name, - asname=import_alias.evaluated_alias, + # Add removal tuple instead of calling directly + self.scheduled_removals.add( + ( + import_alias.evaluated_name, + None, + import_alias.evaluated_alias, + ) ) new_names.append(import_alias.with_changes(asname=None)) @@ -207,7 +214,7 @@ def leave_ImportFrom( return updated_node else: - new_names = [] + new_names: list[cst.ImportAlias] = [] for import_alias in names: alias_name = get_full_name_for_node(import_alias.name) if alias_name is not None: @@ -245,6 +252,10 @@ def leave_ImportFrom( # This import might be in use elsewhere in the code, so schedule a potential removal. self.scheduled_removals.add(original_node) new_names.append(import_alias) + if isinstance(new_names[-1].comma, cst.Comma): + new_names[-1] = new_names[-1].with_changes( + comma=cst.MaybeSentinel.DEFAULT + ) return updated_node.with_changes(names=new_names) return updated_node @@ -307,10 +318,13 @@ def leave_Attribute( def leave_Module( self, original_node: cst.Module, updated_node: cst.Module ) -> cst.Module: - for removal_node in self.scheduled_removals: - RemoveImportsVisitor.remove_unused_import_by_node( - self.context, removal_node - ) + for removal in self.scheduled_removals: + if isinstance(removal, tuple): + RemoveImportsVisitor.remove_unused_import( + self.context, removal[0], removal[1], removal[2] + ) + else: + RemoveImportsVisitor.remove_unused_import_by_node(self.context, removal) # If bypass_import is False, we know that no import statements were directly renamed, and the fact # that we have any `self.scheduled_removals` tells us we encountered a matching `old_name` in the code. if not self.bypass_import and self.scheduled_removals: diff --git a/libcst/codemod/commands/tests/test_rename.py b/libcst/codemod/commands/tests/test_rename.py index 20e1c7d4..8245a34c 100644 --- a/libcst/codemod/commands/tests/test_rename.py +++ b/libcst/codemod/commands/tests/test_rename.py @@ -382,6 +382,28 @@ class Foo(d.z): new_name="d.z", ) + def test_comma_import(self) -> None: + before = """ + import a, b, c + + class Foo(a.z): + bar: b.bar + baz: c.baz + """ + after = """ + import a, b, d + + class Foo(a.z): + bar: b.bar + baz: d.baz + """ + self.assertCodemod( + before, + after, + old_name="c.baz", + new_name="d.baz", + ) + def test_other_import_froms_untouched(self) -> None: before = """ from a import b, c, d @@ -405,6 +427,29 @@ class Foo(b): new_name="f.b", ) + def test_comma_import_from(self) -> None: + before = """ + from a import b, c, d + + class Foo(b): + bar: c.bar + baz: d.baz + """ + after = """ + from a import b, c + from f import d + + class Foo(b): + bar: c.bar + baz: d.baz + """ + self.assertCodemod( + before, + after, + old_name="a.d", + new_name="f.d", + ) + def test_no_removal_of_import_in_use(self) -> None: before = """ import a