From a36432c958fc34bdf28e19f2962ac9f70d35d9f9 Mon Sep 17 00:00:00 2001
From: Zsolt Dollenstein <zsol.zsol@gmail.com>
Date: Fri, 29 Nov 2024 11:23:59 +0000
Subject: [PATCH] rename: Fix imports with aliases (#1252)

When renaming `a.b` -> `c.d`, in imports like `import a.b as x` the as_name wasn't correctly removed even though references to `x` were renamed to `c.d`.

This PR makes the codemod remove the `x` asname in these cases.
---
 libcst/codemod/commands/rename.py            | 16 ++++++++--
 libcst/codemod/commands/tests/test_rename.py | 32 ++++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

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")