diff --git a/docs/source/codemods.rst b/docs/source/codemods.rst index aa59be1f..4e563472 100644 --- a/docs/source/codemods.rst +++ b/docs/source/codemods.rst @@ -151,3 +151,5 @@ inside codemods. As of now, the list includes the following helpers. :exclude-members: visit_AnnAssign, leave_AnnAssign, visit_Assign, leave_Assign, visit_List, leave_List, visit_Tuple, leave_Tuple, visit_Set, leave_Set, visit_Element .. autoclass:: libcst.codemod.visitors.AddImportsVisitor :exclude-members: CONTEXT_KEY, visit_Module, leave_ImportFrom, leave_Module +.. autoclass:: libcst.codemod.visitors.RemoveImportsVisitor + :exclude-members: CONTEXT_KEY, METADATA_DEPENDENCIES, visit_Module, leave_ImportFrom, leave_Import diff --git a/libcst/codemod/_command.py b/libcst/codemod/_command.py index 1774558a..f9374b03 100644 --- a/libcst/codemod/_command.py +++ b/libcst/codemod/_command.py @@ -14,6 +14,7 @@ from libcst.codemod._codemod import Codemod from libcst.codemod._context import CodemodContext from libcst.codemod._visitor import ContextAwareTransformer from libcst.codemod.visitors._add_imports import AddImportsVisitor +from libcst.codemod.visitors._remove_imports import RemoveImportsVisitor _Codemod = TypeVar("_Codemod", bound=Codemod) @@ -31,7 +32,8 @@ class CodemodCommand(Codemod, ABC): The following list of transforms are automatically run at this time: - - :class:`~libcst.codemod.visitors.AddImportsVisitor` (adds needed imports to a file). + - :class:`~libcst.codemod.visitors.AddImportsVisitor` (adds needed imports to a module). + - :class:`~libcst.codemod.visitors.RemoveImportsVisitor` (removes unreferenced imports from a module). """ #: An overrideable description attribute so that codemods can provide @@ -76,7 +78,8 @@ class CodemodCommand(Codemod, ABC): # a context and other optional args and modifies its own context key # accordingly. We import them here so that we don't have circular imports. supported_transforms: Dict[str, Type[Codemod]] = { - AddImportsVisitor.CONTEXT_KEY: AddImportsVisitor + AddImportsVisitor.CONTEXT_KEY: AddImportsVisitor, + RemoveImportsVisitor.CONTEXT_KEY: RemoveImportsVisitor, } # For any visitors that we support auto-running, run them here if needed. diff --git a/libcst/codemod/visitors/__init__.py b/libcst/codemod/visitors/__init__.py index 2e66c8fa..ac3bec55 100644 --- a/libcst/codemod/visitors/__init__.py +++ b/libcst/codemod/visitors/__init__.py @@ -7,6 +7,12 @@ from libcst.codemod.visitors._add_imports import AddImportsVisitor from libcst.codemod.visitors._gather_exports import GatherExportsVisitor from libcst.codemod.visitors._gather_imports import GatherImportsVisitor +from libcst.codemod.visitors._remove_imports import RemoveImportsVisitor -__all__ = ["AddImportsVisitor", "GatherImportsVisitor", "GatherExportsVisitor"] +__all__ = [ + "AddImportsVisitor", + "GatherImportsVisitor", + "GatherExportsVisitor", + "RemoveImportsVisitor", +] diff --git a/libcst/codemod/visitors/_add_imports.py b/libcst/codemod/visitors/_add_imports.py index d79c9600..272edfc7 100644 --- a/libcst/codemod/visitors/_add_imports.py +++ b/libcst/codemod/visitors/_add_imports.py @@ -19,9 +19,10 @@ class AddImportsVisitor(ContextAwareTransformer): """ Ensures that given imports exist in a module. Given a :class:`~libcst.codemod.CodemodContext` and a sequence of tuples specifying - a module to import from as a string and optionally an object to import from - that module, ensures that that import exists. It will modify existing imports - as necessary if the module in question is already being imported from. + a module to import from as a string. Optionally an object to import from + that module and any alias to assign that import, ensures that that + import exists. It will modify existing imports as necessary if the module + in question is already being imported from. This is one of the transforms that is available automatically to you when running a codemod. To use it in this manner, import @@ -29,7 +30,8 @@ class AddImportsVisitor(ContextAwareTransformer): :meth:`~libcst.codemod.visitors.AddImportsVisitor.add_needed_import` method, giving it the current context (found as ``self.context`` for all subclasses of :class:`~libcst.codemod.Codemod`), the module you wish to import from and - optionally an object you wish to import from that module. + optionally an object you wish to import from that module and any alias you + would like to assign that import to. For example:: @@ -78,12 +80,18 @@ class AddImportsVisitor(ContextAwareTransformer): """ Schedule an import to be added in a future invocation of this class by updating the ``context`` to include the ``module`` and optionally ``obj`` - to be imported. When subclassing from + to be imported as well as optionally ``alias`` to alias the imported + ``module`` or ``obj`` to. 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.CSTNode.visit` method on the module under modification - with an instance of this class after performing your transform. + :meth:`~libcst.codemod.Codemod.transform_module` method on the module + under modification with an instance of this class after performing your + transform. Note that if the particular ``module`` or ``obj`` you are + requesting to import already exists as an import on the current module + at the time of executing :meth:`~libcst.codemod.Codemod.transform_module` + on an instance of :class:`~libcst.codemod.visitors.AddImportsVisitor`, + this will perform no action in order to avoid adding duplicate imports. """ if module == "__future__" and obj is None: diff --git a/libcst/codemod/visitors/_remove_imports.py b/libcst/codemod/visitors/_remove_imports.py new file mode 100644 index 00000000..d1cc8b8c --- /dev/null +++ b/libcst/codemod/visitors/_remove_imports.py @@ -0,0 +1,421 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# +# pyre-strict +from typing import Dict, List, Optional, Sequence, Set, Tuple, Union + +import libcst as cst +from libcst.codemod._context import CodemodContext +from libcst.codemod._visitor import ContextAwareTransformer, ContextAwareVisitor +from libcst.codemod.visitors._gather_exports import GatherExportsVisitor +from libcst.helpers import get_absolute_module_for_import, get_full_name_for_node +from libcst.metadata import Assignment, Scope, ScopeProvider + + +class RemovedNodeVisitor(ContextAwareVisitor): + def _remove_imports_from_import_stmt( + self, local_name: str, import_node: cst.Import + ) -> None: + for import_alias in import_node.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, + import_alias.evaluated_name, + asname=import_alias.evaluated_alias, + ) + + 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 _visit_name_attr_alike(self, node: Union[cst.Name, cst.Attribute]) -> None: + # Look up the local name of this node. + local_name = get_full_name_for_node(node) + if local_name is None: + return + + # Look up the scope for this node, remove the import that caused it to exist. + metadata_wrapper = self.context.wrapper + if metadata_wrapper is None: + raise Exception("Cannot look up import, metadata is not computed for node!") + scope_provider = metadata_wrapper.resolve(ScopeProvider) + try: + scope = scope_provider[node] + if scope is None: + # This object has no scope, so we can't remove it. + return + except KeyError: + # This object has no scope, so we can't remove it. + return + + while True: + for assignment in scope.assignments[node] or set(): + # We only care about non-builtins. + if isinstance(assignment, Assignment): + import_node = assignment.node + if isinstance(import_node, cst.Import): + self._remove_imports_from_import_stmt(local_name, import_node) + elif isinstance(import_node, cst.ImportFrom): + self._remove_imports_from_importfrom_stmt( + local_name, import_node + ) + + if scope is scope.parent: + break + scope = scope.parent + + def visit_Name(self, node: cst.Name) -> None: + self._visit_name_attr_alike(node) + + def visit_Attribute(self, node: cst.Attribute) -> None: + self._visit_name_attr_alike(node) + + +class RemoveImportsVisitor(ContextAwareTransformer): + """ + Attempt to remove given imports from a module, dependent on whether there are + any uses of the imported objects. Given a :class:`~libcst.codemod.CodemodContext` + and a sequence of tuples specifying a module to remove as a string. Optionally + an object being imported from that module and optionally an alias assigned to + that imported object, ensures that that import no longer exists as long as there + are no remaining references. + + Note that static analysis is able to determine safely whether an import is still + needed given a particular module, but it is currently unable to determine whether + an imported object is re-exported and used inside another module unless that + object appears in an ``__any__`` list. + + This is one of the transforms that is available automatically to you when running + a codemod. To use it in this manner, importi + :class:`~libcst.codemod.visitors.RemoveImportsVisitor` and then call the static + :meth:`~libcst.codemod.visitors.RemoveImportsVisitor.remove_unused_import` method, + giving it the current context (found as ``self.context`` for all subclasses of + :class:`~libcst.codemod.Codemod`), the module you wish to remove and + optionally an object you wish to stop importing as well as an alias that the + object is currently assigned to. + + For example:: + + RemoveImportsVisitor.remove_unused_import(self.context, "typing", "Optional") + + This will remove any ``from typing import Optional`` that exists in the module + as long as there are no uses of ``Optional`` in that module. + + As another example:: + + RemoveImportsVisitor.remove_unused_import(self.context, "typing") + + This will remove any ``import typing`` that exists in the module, as long as + there are no references to ``typing`` in that module, including references + such as ``typing.Optional``. + + Additionally, :class:`~libcst.codemod.visitors.RemoveImportsVisitor` includes + a convenience function + :meth:`~libcst.codemod.visitors.RemoveImportsVisitor.remove_unused_import_by_node` + which will attempt to schedule removal of all imports referenced in that node + and its children. This is especially useful inside transforms when you are going + to remove a node using :func:`~libcst.RemoveFromParent` to get rid of a node. + + For example:: + + def leave_AnnAssign( + self, original_node: cst.AnnAssign, updated_node: cst.AnnAssign, + ) -> cst.RemovalSentinel: + # Remove all annotated assignment statements, clean up imports. + RemoveImportsVisitor.remove_unused_import_by_node(self.context, original_node) + return cst.RemovalFromParent() + + This will remove all annotated assignment statements from a module as well + as clean up any imports that were only referenced in those assignments. Note + that we pass the ``original_node`` to the helper function as it uses scope analysis + under the hood which is only computed on the original tree. + + Note that this is a subclass of :class:`~libcst.CSTTransformer` so it is + possible to instantiate it and pass it to a :class:`~libcst.Module` + :meth:`~libcst.CSTNode.visit` method. However, it is far easier to use + the automatic transform feature of :class:`~libcst.codemod.CodemodCommand` + and schedule an import to be added by calling + :meth:`~libcst.codemod.visitors.RemoveImportsVisitor.remove_unused_import` + + """ + + CONTEXT_KEY = "RemoveImportsVisitor" + METADATA_DEPENDENCIES = (ScopeProvider,) + + @staticmethod + def _get_imports_from_context( + context: CodemodContext, + ) -> List[Tuple[str, Optional[str], Optional[str]]]: + unused_imports = context.scratch.get(RemoveImportsVisitor.CONTEXT_KEY, []) + if not isinstance(unused_imports, list): + raise Exception("Logic error!") + return unused_imports + + @staticmethod + def remove_unused_import( + context: CodemodContext, + module: str, + obj: Optional[str] = None, + asname: Optional[str] = None, + ) -> None: + """ + Schedule an import to be removed in a future invocation of this class by + updating the ``context`` to include the ``module`` and optionally ``obj`` + which is currently imported as well as optionally ``alias`` that the + imported ``module`` or ``obj`` is aliased to. 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 if the particular ``module`` or ``obj`` you are + requesting to remove is still in use somewhere in the current module + at the time of executing :meth:`~libcst.codemod.Codemod.transform_module` + on an instance of :class:`~libcst.codemod.visitors.AddImportsVisitor`, + this will perform no action in order to avoid removing an in-use import. + """ + + unused_imports = RemoveImportsVisitor._get_imports_from_context(context) + unused_imports.append((module, obj, asname)) + context.scratch[RemoveImportsVisitor.CONTEXT_KEY] = unused_imports + + @staticmethod + 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 __init__( + self, + context: CodemodContext, + unused_imports: Sequence[Tuple[str, Optional[str], Optional[str]]] = (), + ) -> None: + # Allow for instantiation from either a context (used when multiple transforms + # get chained) or from a direct instantiation. + super().__init__(context) + + all_unused_imports: List[Tuple[str, Optional[str], Optional[str]]] = [ + *RemoveImportsVisitor._get_imports_from_context(context), + *unused_imports, + ] + self.unused_module_imports: Dict[str, Optional[str]] = { + module: alias for module, obj, alias in all_unused_imports if obj is None + } + self.unused_obj_imports: Dict[str, Set[Tuple[str, Optional[str]]]] = {} + self.exported_objects: Set[str] = set() + for module, obj, alias in all_unused_imports: + if obj is None: + continue + if module not in self.unused_obj_imports: + self.unused_obj_imports[module] = set() + self.unused_obj_imports[module].add((obj, alias)) + + def visit_Module(self, node: cst.Module) -> None: + object_visitor = GatherExportsVisitor(self.context) + node.visit(object_visitor) + self.exported_objects = object_visitor.explicit_exported_objects + + def _is_in_use(self, scope: Scope, alias: cst.ImportAlias) -> bool: + # Grab the string name of this alias from the point of view of this module. + asname = alias.asname + if asname is not None: + name_node = asname.name + else: + name_node = alias.name + while isinstance(name_node, cst.Attribute): + name_node = name_node.value + name_or_alias = cst.ensure_type(name_node, cst.Name).value + + 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) + ): + 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) + + def leave_Import( + self, original_node: cst.Import, updated_node: cst.Import + ) -> Union[cst.Import, 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 + + names_to_keep = [] + for import_alias in original_node.names: + if import_alias.evaluated_name not in self.unused_module_imports: + # This is a keeper since we aren't removing it + names_to_keep.append(import_alias) + continue + + if ( + import_alias.evaluated_alias + != self.unused_module_imports[import_alias.evaluated_name] + ): + # This is a keeper since the alias does not match + # what we are looking for. + names_to_keep.append(import_alias) + continue + + # Now that we know we want to remove this module, 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 + + # Now, either remove this statement or remove the imports we are + # deleting from this statement. + if len(names_to_keep) == 0: + return cst.RemoveFromParent() + else: + # 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 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 + + # Now, either remove this statement or remove the imports we are + # deleting from this statement. + if len(names_to_keep) == 0: + return cst.RemoveFromParent() + else: + # 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) diff --git a/libcst/codemod/visitors/tests/test_remove_imports.py b/libcst/codemod/visitors/tests/test_remove_imports.py new file mode 100644 index 00000000..4e2801c2 --- /dev/null +++ b/libcst/codemod/visitors/tests/test_remove_imports.py @@ -0,0 +1,756 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# +# pyre-strict +import libcst as cst +import libcst.matchers as m +from libcst.codemod import CodemodContext, CodemodTest, VisitorBasedCodemodCommand +from libcst.codemod.visitors import RemoveImportsVisitor +from libcst.metadata import ( + QualifiedName, + QualifiedNameProvider, + QualifiedNameSource, + ScopeProvider, +) +from libcst.testing.utils import data_provider + + +class TestRemoveImportsCodemod(CodemodTest): + + TRANSFORM = RemoveImportsVisitor + + def test_noop(self) -> None: + """ + Should do nothing. + """ + + before = """ + def foo() -> None: + pass + """ + after = """ + def foo() -> None: + pass + """ + + self.assertCodemod(before, after, []) + + def test_remove_import_simple(self) -> None: + """ + Should remove module as import + """ + + before = """ + import bar + import baz + + def foo() -> None: + pass + """ + after = """ + import bar + + def foo() -> None: + pass + """ + + self.assertCodemod(before, after, [("baz", None, None)]) + + def test_remove_import_alias_simple(self) -> None: + """ + Should remove aliased module as import + """ + + before = """ + import bar + import baz as qux + + def foo() -> None: + pass + """ + after = """ + import bar + + def foo() -> None: + pass + """ + + self.assertCodemod(before, after, [("baz", None, "qux")]) + + def test_dont_remove_import_simple(self) -> None: + """ + Should not remove module import with reference + """ + + before = """ + import bar + import baz + + def foo() -> None: + baz.qux() + """ + after = """ + import bar + import baz + + def foo() -> None: + baz.qux() + """ + + self.assertCodemod(before, after, [("baz", None, None)]) + + def test_dont_remove_import_alias_simple(self) -> None: + """ + Should not remove aliased module import with reference + """ + + before = """ + import bar + import baz as qux + + def foo() -> None: + qux.quux() + """ + after = """ + import bar + import baz as qux + + def foo() -> None: + qux.quux() + """ + + self.assertCodemod(before, after, [("baz", None, "qux")]) + + def test_dont_remove_import_simple_wrong_alias(self) -> None: + """ + Should not remove module as import since wrong alias + """ + + before = """ + import bar + import baz + + def foo() -> None: + pass + """ + after = """ + import bar + import baz + + def foo() -> None: + pass + """ + + self.assertCodemod(before, after, [("baz", None, "qux")]) + + def test_dont_remove_import_wrong_alias_simple(self) -> None: + """ + Should not remove wrong aliased module as import + """ + + before = """ + import bar + import baz as qux + + def foo() -> None: + pass + """ + after = """ + import bar + import baz as qux + + def foo() -> None: + pass + """ + + self.assertCodemod(before, after, [("baz", None, None)]) + + def test_remove_importfrom_simple(self) -> None: + """ + Should remove import from + """ + + before = """ + import bar + from baz import qux + + def foo() -> None: + pass + """ + after = """ + import bar + + def foo() -> None: + pass + """ + + self.assertCodemod(before, after, [("baz", "qux", None)]) + + def test_remove_importfrom_alias_simple(self) -> None: + """ + Should remove import from with alias + """ + + before = """ + import bar + from baz import qux as quux + + def foo() -> None: + pass + """ + after = """ + import bar + + def foo() -> None: + pass + """ + + self.assertCodemod(before, after, [("baz", "qux", "quux")]) + + def test_dont_remove_importfrom_simple(self) -> None: + """ + Should not remove import from with reference + """ + + before = """ + import bar + from baz import qux + + def foo() -> None: + qux() + """ + after = """ + import bar + from baz import qux + + def foo() -> None: + qux() + """ + + self.assertCodemod(before, after, [("baz", "qux", None)]) + + def test_dont_remove_importfrom_alias_simple(self) -> None: + """ + Should not remove aliased import from with reference + """ + + before = """ + import bar + from baz import qux as quux + + def foo() -> None: + quux() + """ + after = """ + import bar + from baz import qux as quux + + def foo() -> None: + quux() + """ + + self.assertCodemod(before, after, [("baz", "qux", "quux")]) + + def test_dont_remove_importfrom_simple_wrong_alias(self) -> None: + """ + Should not remove import from since it is wrong alias + """ + + before = """ + import bar + from baz import qux as quux + + def foo() -> None: + pass + """ + after = """ + import bar + from baz import qux as quux + + def foo() -> None: + pass + """ + + self.assertCodemod(before, after, [("baz", "qux", None)]) + + def test_dont_remove_importfrom_alias_simple_wrong_alias(self) -> None: + """ + Should not remove import from with wrong alias + """ + + before = """ + import bar + from baz import qux + + def foo() -> None: + pass + """ + after = """ + import bar + from baz import qux + + def foo() -> None: + pass + """ + + self.assertCodemod(before, after, [("baz", "qux", "quux")]) + + def test_remove_importfrom_relative(self) -> None: + """ + Should remove import from which is relative + """ + + before = """ + import bar + from .c import qux + + def foo() -> None: + pass + """ + after = """ + import bar + + def foo() -> None: + pass + """ + + self.assertCodemod( + before, + after, + [("a.b.c", "qux", None)], + context_override=CodemodContext(full_module_name="a.b.foobar"), + ) + + def test_dont_remove_inuse_importfrom_relative(self) -> None: + """ + Should not remove import from which is relative since it is in use. + """ + + before = """ + import bar + from .c import qux + + def foo() -> None: + qux() + """ + after = """ + import bar + from .c import qux + + def foo() -> None: + qux() + """ + + self.assertCodemod( + before, + after, + [("a.b.c", "qux", None)], + context_override=CodemodContext(full_module_name="a.b.foobar"), + ) + + def test_dont_remove_wrong_importfrom_relative(self) -> None: + """ + Should not remove import from which is relative since it is the wrong module. + """ + + before = """ + import bar + from .c import qux + + def foo() -> None: + pass + """ + after = """ + import bar + from .c import qux + + def foo() -> None: + pass + """ + + self.assertCodemod( + before, + after, + [("a.b.d", "qux", None)], + context_override=CodemodContext(full_module_name="a.b.foobar"), + ) + + def test_remove_import_complex(self) -> None: + """ + Should remove complex module as import + """ + + before = """ + import bar + import baz, qux + import a.b + import c.d + import e.f as g + import h.i as j + + def foo() -> None: + c.d() + j() + """ + after = """ + import bar + import qux + import c.d + import h.i as j + + def foo() -> None: + c.d() + j() + """ + + self.assertCodemod( + before, + after, + [ + ("baz", None, None), + ("a.b", None, None), + ("c.d", None, None), + ("e.f", None, "g"), + ("h.i", None, "j"), + ], + ) + + def test_remove_fromimport_complex(self) -> None: + """ + Should remove complex from import + """ + + before = """ + from bar import qux, quux + from a.b import c + from d.e import f + from h.i import j as k + from l.m import n as o + + def foo() -> None: + f() + k() + """ + after = """ + from bar import qux + from d.e import f + from h.i import j as k + + def foo() -> None: + f() + k() + """ + + self.assertCodemod( + before, + after, + [ + ("bar", "quux", None), + ("a.b", "c", None), + ("d.e", "f", None), + ("h.i", "j", "k"), + ("l.m", "n", "o"), + ], + ) + + def test_remove_import_multiple_assignments(self) -> None: + """ + Should remove import with multiple assignments + """ + + before = """ + from foo import bar + from qux import bar + + def foo() -> None: + bar() + """ + after = """ + from qux import bar + + def foo() -> None: + bar() + """ + + self.assertCodemod(before, after, [("foo", "bar", None)]) + + @data_provider( + ( + # Simple removal, no other uses. + ( + """ + from foo import bar + from qux import baz + + def fun() -> None: + bar() + baz() + """, + """ + from qux import baz + + def fun() -> None: + baz() + """, + ), + # Remove a node, other uses, don't remove import. + ( + """ + from foo import bar + from qux import baz + + def fun() -> None: + bar() + baz() + + def foobar() -> None: + a = bar + a() + """, + """ + from foo import bar + from qux import baz + + def fun() -> None: + baz() + + def foobar() -> None: + a = bar + a() + """, + ), + # Remove an alias. + ( + """ + from foo import bar as other + from qux import baz + + def fun() -> None: + other() + baz() + """, + """ + from qux import baz + + def fun() -> None: + baz() + """, + ), + # Simple removal, no other uses. + ( + """ + import foo + from qux import baz + + def fun() -> None: + foo.bar() + baz() + """, + """ + from qux import baz + + def fun() -> None: + baz() + """, + ), + # Remove a node, other uses, don't remove import. + ( + """ + import foo + from qux import baz + + def fun() -> None: + foo.bar() + baz() + + def foobar() -> None: + a = foo.bar + a() + """, + """ + import foo + from qux import baz + + def fun() -> None: + baz() + + def foobar() -> None: + a = foo.bar + a() + """, + ), + # Remove an alias. + ( + """ + import foo as other + from qux import baz + + def fun() -> None: + other.bar() + baz() + """, + """ + from qux import baz + + def fun() -> None: + baz() + """, + ), + ) + ) + def test_remove_import_by_node_simple(self, before: str, after: str) -> None: + """ + Given a node that's directly referenced in an import, + make sure that the import is removed when the node + is also removed. + """ + + class RemoveBarTransformer(VisitorBasedCodemodCommand): + + METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider) + + @m.leave( + m.SimpleStatementLine( + body=[ + m.Expr( + m.Call( + metadata=m.MatchMetadata( + QualifiedNameProvider, + { + QualifiedName( + source=QualifiedNameSource.IMPORT, + name="foo.bar", + ) + }, + ) + ) + ) + ] + ) + ) + def _leave_foo_bar( + self, + original_node: cst.SimpleStatementLine, + updated_node: cst.SimpleStatementLine, + ) -> cst.RemovalSentinel: + RemoveImportsVisitor.remove_unused_import_by_node( + self.context, original_node + ) + return cst.RemoveFromParent() + + module = cst.parse_module(self.make_fixture_data(before)) + self.assertCodeEqual( + after, RemoveBarTransformer(CodemodContext()).transform_module(module).code + ) + + def test_remove_import_from_node(self) -> None: + """ + Make sure that if an import node itself is requested for + removal, we still do the right thing and only remove it + if it is unused. + """ + + before = """ + from foo import bar + from qux import baz + from foo import qux as other + from qux import foobar as other2 + + def fun() -> None: + baz() + other2() + """ + after = """ + from qux import baz + from qux import foobar as other2 + + def fun() -> None: + baz() + other2() + """ + + class RemoveImportTransformer(VisitorBasedCodemodCommand): + + METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider) + + def visit_ImportFrom(self, node: cst.ImportFrom) -> None: + RemoveImportsVisitor.remove_unused_import_by_node(self.context, node) + + module = cst.parse_module(self.make_fixture_data(before)) + self.assertCodeEqual( + after, + RemoveImportTransformer(CodemodContext()).transform_module(module).code, + ) + + def test_remove_import_node(self) -> None: + """ + Make sure that if an import node itself is requested for + removal, we still do the right thing and only remove it + if it is unused. + """ + + before = """ + import foo + import qux + import bar as other + import foobar as other2 + + def fun() -> None: + qux.baz() + other2.baz() + """ + after = """ + import qux + import foobar as other2 + + def fun() -> None: + qux.baz() + other2.baz() + """ + + class RemoveImportTransformer(VisitorBasedCodemodCommand): + + METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider) + + def visit_Import(self, node: cst.Import) -> None: + RemoveImportsVisitor.remove_unused_import_by_node(self.context, node) + + module = cst.parse_module(self.make_fixture_data(before)) + self.assertCodeEqual( + after, + RemoveImportTransformer(CodemodContext()).transform_module(module).code, + ) + + def test_remove_import_with_all(self) -> None: + """ + Make sure that if an import node itself is requested for + removal, we don't remove it if it shows up in an __all__ + node. + """ + + before = """ + from foo import bar + from qux import baz + + __all__ = ["baz"] + """ + after = """ + from qux import baz + + __all__ = ["baz"] + """ + + class RemoveImportTransformer(VisitorBasedCodemodCommand): + + METADATA_DEPENDENCIES = (QualifiedNameProvider, ScopeProvider) + + def visit_ImportFrom(self, node: cst.ImportFrom) -> None: + RemoveImportsVisitor.remove_unused_import_by_node(self.context, node) + + module = cst.parse_module(self.make_fixture_data(before)) + self.assertCodeEqual( + after, + RemoveImportTransformer(CodemodContext()).transform_module(module).code, + )