Skip to content

Commit

Permalink
rename: Fix imports with aliases
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
zsol committed Nov 29, 2024
1 parent 28e0f39 commit 4106e2a
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 3 deletions.
16 changes: 13 additions & 3 deletions libcst/codemod/commands/rename.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)

Expand Down
32 changes: 32 additions & 0 deletions libcst/codemod/commands/tests/test_rename.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")

0 comments on commit 4106e2a

Please sign in to comment.