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)