def leave_ImportFrom(self, original_node: libcst.ImportFrom,
updated_node: libcst.ImportFrom) -> libcst.ImportFrom:
if isinstance(updated_node.names, libcst.ImportStar):
# There's nothing to do here!
return updated_node
# Get the module we're importing as a string, see if we have work to do.
module = get_absolute_module_for_import(self.context.full_module_name,
updated_node)
if (module is None or module not in self.module_mapping
and module not in self.alias_mapping):
return updated_node
# We have work to do, mark that we won't modify this again.
imports_to_add = self.module_mapping.get(module, [])
if module in self.module_mapping:
del self.module_mapping[module]
aliases_to_add = self.alias_mapping.get(module, [])
if module in self.alias_mapping:
del self.alias_mapping[module]
# Now, do the actual update.
return updated_node.with_changes(names=[
*(libcst.ImportAlias(name=libcst.Name(imp))
for imp in sorted(imports_to_add)),
*(libcst.ImportAlias(
name=libcst.Name(imp),
asname=libcst.AsName(name=libcst.Name(alias)),
) for (imp, alias) in sorted(aliases_to_add)),
*updated_node.names,
])
def _remove_imports_from_importfrom_stmt(
self, local_name: str, import_node: cst.ImportFrom) -> None:
names = import_node.names
if isinstance(names, cst.ImportStar):
# We don't handle removing this, so ignore it.
return
module_name = get_absolute_module_for_import(
self.context.full_module_name, import_node)
if module_name is None:
raise Exception(
"Cannot look up absolute module from relative import!")
# We know any local names will refer to this as an alias if
# there is one, and as the original name if there is not one
for import_alias in names:
if import_alias.evaluated_alias is None:
prefix = import_alias.evaluated_name
else:
prefix = import_alias.evaluated_alias
if local_name == prefix or local_name.startswith(f"{prefix}."):
RemoveImportsVisitor.remove_unused_import(
self.context,
module_name,
obj=import_alias.evaluated_name,
asname=import_alias.evaluated_alias,
)
def _handle_import(self, node: Union[Import, ImportFrom]) -> None:
node_start = self.get_metadata(PositionProvider, node).start.line
if node_start in self._ignored_lines:
return
names = node.names
if isinstance(names, ImportStar):
return
for alias in names:
position = self.get_metadata(PositionProvider, alias)
lines = set(range(position.start.line, position.end.line + 1))
if lines.isdisjoint(self._ignored_lines):
if isinstance(node, Import):
RemoveImportsVisitor.remove_unused_import(
self.context,
module=alias.evaluated_name,
asname=alias.evaluated_alias,
)
else:
module_name = get_absolute_module_for_import(
self.context.full_module_name, node)
if module_name is None:
raise ValueError(
f"Couldn't get absolute module name for {alias.evaluated_name}"
)
RemoveImportsVisitor.remove_unused_import(
self.context,
module=module_name,
obj=alias.evaluated_name,
asname=alias.evaluated_alias,
)
def leave_ImportFrom(
self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
) -> Union[cst.ImportFrom, cst.RemovalSentinel]:
# Grab the scope for this import. If we don't have scope, we can't determine
# whether this import is unused so it is unsafe to remove.
scope = self.get_metadata(ScopeProvider, original_node, None)
if scope is None:
return updated_node
# Make sure we have anything to do with this node.
names = original_node.names
if isinstance(names, cst.ImportStar):
# This is a star import, so we won't remove it.
return updated_node
# Make sure we actually know the absolute module.
module_name = get_absolute_module_for_import(
self.context.full_module_name, updated_node
)
if module_name is None or module_name not in self.unused_obj_imports:
# This node isn't on our list of todos, so let's bail.
return updated_node
objects_to_remove = self.unused_obj_imports[module_name]
names_to_keep = []
for import_alias in names:
# Figure out if it is in our list of things to kill
for name, alias in objects_to_remove:
if (
name == import_alias.evaluated_name
and alias == import_alias.evaluated_alias
):
break
else:
# This is a keeper, we don't have it on our list.
names_to_keep.append(import_alias)
continue
# Now that we know we want to remove this object, figure out if
# there are any live references to it.
if self._is_in_use(scope, import_alias):
names_to_keep.append(import_alias)
continue
# no changes
if names_to_keep == names:
return updated_node
# Now, either remove this statement or remove the imports we are
# deleting from this statement.
if len(names_to_keep) == 0:
return cst.RemoveFromParent()
if names_to_keep[-1] != names[-1]:
# Remove trailing comma in order to not mess up import statements.
names_to_keep = [
*names_to_keep[:-1],
names_to_keep[-1].with_changes(comma=cst.MaybeSentinel.DEFAULT),
]
return updated_node.with_changes(names=names_to_keep)
def remove_unused_import_by_node(context: CodemodContext,
node: cst.CSTNode) -> None:
"""
Schedule any imports referenced by ``node`` or one of its children
to be removed in a future invocation of this class by updating the
``context`` to include the ``module``, ``obj`` and ``alias`` for each
import in question. When subclassing from
:class:`~libcst.codemod.CodemodCommand`, this will be performed for you
after your transform finishes executing. If you are subclassing from a
:class:`~libcst.codemod.Codemod` instead, you will need to call the
:meth:`~libcst.codemod.Codemod.transform_module` method on the module
under modification with an instance of this class after performing your
transform. Note that all imports that are referenced by this ``node``
or its children will only be removed if they are not in use at the time
of exeucting :meth:`~libcst.codemod.Codemod.transform_module`
on an instance of :class:`~libcst.codemod.visitors.AddImportsVisitor`
in order to avoid removing an in-use import.
"""
# Special case both Import and ImportFrom so they can be
# directly removed here.
if isinstance(node, cst.Import):
for import_alias in node.names:
RemoveImportsVisitor.remove_unused_import(
context,
import_alias.evaluated_name,
asname=import_alias.evaluated_alias,
)
elif isinstance(node, cst.ImportFrom):
names = node.names
if isinstance(names, cst.ImportStar):
# We don't handle removing this, so ignore it.
return
module_name = get_absolute_module_for_import(
context.full_module_name, node)
if module_name is None:
raise Exception(
"Cannot look up absolute module from relative import!")
for import_alias in names:
RemoveImportsVisitor.remove_unused_import(
context,
module_name,
obj=import_alias.evaluated_name,
asname=import_alias.evaluated_alias,
)
else:
# Look up all children that could have been imported. Any that
# we find will be scheduled for removal.
node.visit(RemovedNodeVisitor(context))
def test_get_absolute_module(
self, module: Optional[str], importfrom: str, output: Optional[str],
) -> None:
node = ensure_type(cst.parse_statement(importfrom), cst.SimpleStatementLine)
assert len(node.body) == 1, "Unexpected number of statements!"
import_node = ensure_type(node.body[0], cst.ImportFrom)
self.assertEqual(get_absolute_module_for_import(module, import_node), output)
if output is None:
with self.assertRaises(Exception):
get_absolute_module_for_import_or_raise(module, import_node)
else:
self.assertEqual(
get_absolute_module_for_import_or_raise(module, import_node), output
)
def visit_ImportFrom(self, node: libcst.ImportFrom) -> None:
# Track this import statement for later analysis.
self.all_imports.append(node)
# Get the module we're importing as a string.
module = get_absolute_module_for_import(self.context.full_module_name,
node)
if module is None:
# Can't get the absolute import from relative, so we can't
# support this.
return
nodenames = node.names
if isinstance(nodenames, libcst.ImportStar):
# We cover everything, no need to bother tracking other things
self.object_mapping[module] = set("*")
return
elif isinstance(nodenames, Sequence):
# Get the list of imports we're aliasing in this import
new_aliases = [(ia.evaluated_name, ia.evaluated_alias)
for ia in nodenames if ia.asname is not None]
if new_aliases:
if module not in self.alias_mapping:
self.alias_mapping[module] = []
# pyre-ignore We know that aliases are not None here.
self.alias_mapping[module].extend(new_aliases)
# Get the list of imports we're importing in this import
new_objects = {
ia.evaluated_name
for ia in nodenames if ia.asname is None
}
if new_objects:
if module not in self.object_mapping:
self.object_mapping[module] = set()
# Make sure that we don't add to a '*' module
if "*" in self.object_mapping[module]:
self.object_mapping[module] = set("*")
return
self.object_mapping[module].update(new_objects)
def leave_ImportFrom(
self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
) -> Union[cst.ImportFrom, cst.RemovalSentinel]:
names = original_node.names
if isinstance(names, cst.ImportStar):
# This is a star import, so we won't remove it.
return updated_node
# Make sure we actually know the absolute module.
module_name = get_absolute_module_for_import(
self.context.full_module_name, updated_node)
if module_name is None or module_name not in self.unused_obj_imports:
# This node isn't on our list of todos, so let's bail.
return updated_node
updates = self._process_importfrom_aliases(updated_node, names,
module_name)
names_to_keep = updates["names"]
# no changes
if names_to_keep == names:
return updated_node
# Now, either remove this statement or remove the imports we are
# deleting from this statement.
if len(names_to_keep) == 0:
return cst.RemoveFromParent()
if names_to_keep[-1] != names[-1]:
# Remove trailing comma in order to not mess up import statements.
names_to_keep = [
*names_to_keep[:-1],
names_to_keep[-1].with_changes(
comma=cst.MaybeSentinel.DEFAULT),
]
updates["names"] = names_to_keep
return updated_node.with_changes(**updates)
