diff --git a/libcst/codemod/commands/rename.py b/libcst/codemod/commands/rename.py index ae7138c8..aad4cea6 100644 --- a/libcst/codemod/commands/rename.py +++ b/libcst/codemod/commands/rename.py @@ -150,9 +150,7 @@ def leave_Import( if import_alias_full_name is None: raise Exception("Could not parse full name for ImportAlias.name node.") - if isinstance( - import_alias_name, (cst.Name, cst.Attribute) - ) and self.old_name.startswith(import_alias_full_name + "."): + if self.old_name.startswith(import_alias_full_name + "."): replacement_module = self.gen_replacement_module(import_alias_full_name) if not replacement_module: # here import_alias_full_name isn't an exact match for old_name @@ -166,6 +164,18 @@ def leave_Import( self.gen_name_or_attr_node(replacement_module) ) new_names.append(cst.ImportAlias(name=new_name_node)) + elif ( + import_alias_full_name == self.new_name + 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, + ) + new_names.append(import_alias.with_changes(asname=None)) return updated_node.with_changes(names=new_names) diff --git a/libcst/codemod/commands/tests/test_rename.py b/libcst/codemod/commands/tests/test_rename.py index efcfbc6e..20e1c7d4 100644 --- a/libcst/codemod/commands/tests/test_rename.py +++ b/libcst/codemod/commands/tests/test_rename.py @@ -111,6 +111,27 @@ def test() -> None: new_name="baz.quux", ) + def test_rename_attr_asname_2(self) -> None: + before = """ + import foo.qux as bar + + def test() -> None: + bar.z(5) + """ + after = """ + import baz.quux + + def test() -> None: + baz.quux.z(5) + """ + + self.assertCodemod( + before, + after, + old_name="foo.qux", + new_name="baz.quux", + ) + def test_rename_module_import(self) -> None: before = """ import a.b @@ -741,3 +762,14 @@ def test_import_parent_module_3(self) -> None: z.c(z.c.d) """ self.assertCodemod(before, after, old_name="a.b.c", new_name="z.c:") + + def test_import_parent_module_asname(self) -> None: + before = """ + import a.b as alias + alias.c(alias.c.d) + """ + after = """ + import z + z.c(z.c.d) + """ + self.assertCodemod(before, after, old_name="a.b.c", new_name="z.c")