Skip to content

Commit

Permalink
RemovedNodeVisitor should not remove imports when references exist.
Browse files Browse the repository at this point in the history
  • Loading branch information
jimmylai committed Mar 16, 2020
1 parent ea8619a commit f612604
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 13 deletions.
18 changes: 6 additions & 12 deletions libcst/codemod/visitors/_remove_imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -306,20 +306,14 @@ def _is_in_use(self, scope: Scope, alias: cst.ImportAlias) -> bool:
if name_or_alias in self.exported_objects:
return True

# number of references to the name
references_count = 0
# number of imports to the same name
assignments_count = 0
for assignment in scope[name_or_alias]:
if isinstance(assignment, Assignment) and isinstance(
assignment.node, (cst.ImportFrom, cst.Import)
if (
isinstance(assignment, Assignment)
and isinstance(assignment.node, (cst.ImportFrom, cst.Import))
and len(assignment.references) > 0
):
assignments_count += 1
references_count += len(assignment.references)

# Remove the import if it's a candidate to remove with no references or
# multiple assignments.
return not (references_count == 0 or assignments_count > 1)
return True
return False

def leave_Import(
self, original_node: cst.Import, updated_node: cst.Import
Expand Down
39 changes: 38 additions & 1 deletion libcst/codemod/visitors/tests/test_remove_imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -457,7 +457,7 @@ def foo() -> None:

def test_remove_import_multiple_assignments(self) -> None:
"""
Should remove import with multiple assignments
Should not remove import with multiple assignments
"""

before = """
Expand All @@ -468,6 +468,7 @@ def foo() -> None:
bar()
"""
after = """
from foo import bar
from qux import bar
def foo() -> None:
Expand All @@ -476,6 +477,42 @@ def foo() -> None:

self.assertCodemod(before, after, [("foo", "bar", None)])

def test_remove_multiple_imports(self) -> None:
"""
Multiple imports
"""
before = """
try:
import a
except Exception:
import a
a.hello()
"""
after = """
try:
import a
except Exception:
import a
a.hello()
"""
self.assertCodemod(before, after, [("a", None, None)])

before = """
try:
import a
except Exception:
import a
"""
after = """
try:
pass
except Exception:
pass
"""
self.assertCodemod(before, after, [("a", None, None)])

@data_provider(
(
# Simple removal, no other uses.
Expand Down

0 comments on commit f612604

Please sign in to comment.