import sys
sys.path.append("../../")
source = """\
import a, b, c as d, e as f # expect to keep: a, c as d
from g import h, i, j as k, l as m # expect to keep: h, j as k
from n import o # expect to be removed entirely
a()
def fun():
d()
class Cls:
att = h.something
def __new__(self) -> "Cls":
var = k.method()
func_undefined(var_undefined)
"""
import libcst as cst
wrapper = cst.metadata.MetadataWrapper(cst.parse_module(source))
scopes = set(wrapper.resolve(cst.metadata.ScopeProvider).values())
for scope in scopes:
print(scope)
from collections import defaultdict
from typing import Dict, Union, Set
unused_imports: Dict[Union[cst.Import, cst.ImportFrom], Set[str]] = defaultdict(set)
undefined_references: Dict[cst.CSTNode, Set[str]] = defaultdict(set)
ranges = wrapper.resolve(cst.metadata.PositionProvider)
for scope in scopes:
for assignment in scope.assignments:
node = assignment.node
if isinstance(assignment, cst.metadata.Assignment) and isinstance(
node, (cst.Import, cst.ImportFrom)
):
if len(assignment.references) == 0:
unused_imports[node].add(assignment.name)
location = ranges[node].start
print(
f"Warning on line {location.line:2d}, column {location.column:2d}: Imported name `{assignment.name}` is unused."
)
for access in scope.accesses:
if len(access.referents) == 0:
node = access.node
location = ranges[node].start
print(
f"Warning on line {location.line:2d}, column {location.column:2d}: Name reference `{node.value}` is not defined."
)
class RemoveUnusedImportTransformer(cst.CSTTransformer):
def __init__(
self, unused_imports: Dict[Union[cst.Import, cst.ImportFrom], Set[str]]
) -> None:
self.unused_imports = unused_imports
def leave_import_alike(
self,
original_node: Union[cst.Import, cst.ImportFrom],
updated_node: Union[cst.Import, cst.ImportFrom],
) -> Union[cst.Import, cst.ImportFrom, cst.RemovalSentinel]:
if original_node not in self.unused_imports:
return updated_node
names_to_keep = []
for name in updated_node.names:
asname = name.asname
if asname is not None:
name_value = asname.name.value
else:
name_value = name.name.value
if name_value not in self.unused_imports[original_node]:
names_to_keep.append(name.with_changes(comma=cst.MaybeSentinel.DEFAULT))
if len(names_to_keep) == 0:
return cst.RemoveFromParent()
else:
return updated_node.with_changes(names=names_to_keep)
def leave_Import(
self, original_node: cst.Import, updated_node: cst.Import
) -> cst.Import:
return self.leave_import_alike(original_node, updated_node)
def leave_ImportFrom(
self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
) -> cst.ImportFrom:
return self.leave_import_alike(original_node, updated_node)
import difflib
fixed_module = wrapper.module.visit(RemoveUnusedImportTransformer(unused_imports))
# Use difflib to show the changes to verify unused imports are removed as expected.
print(
"".join(
difflib.unified_diff(source.splitlines(1), fixed_module.code.splitlines(1))
)
)