[RemoveUnusedImports] Support string type annotations

libcst/codemod/tests/codemod_formatter_error_input.py.txt

@@ -5,7 +5,7 @@
 #
 # pyre-strict
 
-import subprocess  # noqa: F401
+import subprocess
 from contextlib import AsyncExitStack
 
 

libcst/codemod/visitors/__init__.py

@@ -5,15 +5,23 @@
 #
 from libcst.codemod.visitors._add_imports import AddImportsVisitor
 from libcst.codemod.visitors._apply_type_annotations import ApplyTypeAnnotationsVisitor
+from libcst.codemod.visitors._gather_comments import GatherCommentsVisitor
 from libcst.codemod.visitors._gather_exports import GatherExportsVisitor
 from libcst.codemod.visitors._gather_imports import GatherImportsVisitor
+from libcst.codemod.visitors._gather_string_annotation_names import (
+    GatherNamesFromStringAnnotationsVisitor,
+)
+from libcst.codemod.visitors._gather_unused_imports import GatherUnusedImportsVisitor
 from libcst.codemod.visitors._remove_imports import RemoveImportsVisitor
 
 
 __all__ = [
     "AddImportsVisitor",
-    "GatherImportsVisitor",
-    "GatherExportsVisitor",
     "ApplyTypeAnnotationsVisitor",
+    "GatherCommentsVisitor",
+    "GatherExportsVisitor",
+    "GatherImportsVisitor",
+    "GatherNamesFromStringAnnotationsVisitor",
+    "GatherUnusedImportsVisitor",
     "RemoveImportsVisitor",
 ]

libcst/codemod/visitors/_gather_comments.py

@@ -0,0 +1,36 @@
+# 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.
+
+import re
+from typing import Dict, Union
+
+import libcst as cst
+import libcst.matchers as m
+from libcst.codemod._context import CodemodContext
+from libcst.codemod._visitor import ContextAwareVisitor
+from libcst.metadata import PositionProvider
+
+
+class GatherCommentsVisitor(ContextAwareVisitor):
+    METADATA_DEPENDENCIES = (PositionProvider,)
+
+    def __init__(self, context: CodemodContext, comment_regex: str) -> None:
+        super().__init__(context)
+
+        self.comments: Dict[int, cst.Comment] = {}
+
+        self._comment_matcher = re.compile(comment_regex)
+
+    @m.visit(m.EmptyLine(comment=m.DoesNotMatch(None)))
+    @m.visit(m.TrailingWhitespace(comment=m.DoesNotMatch(None)))
+    def visit_comment(self, node: Union[cst.EmptyLine, cst.TrailingWhitespace]) -> None:
+        assert node.comment is not None  # hello, type checker
+        if not self._comment_matcher.match(node.comment.value):
+            return
+        line = self.get_metadata(PositionProvider, node.comment).start.line
+        if isinstance(node, cst.EmptyLine):
+            # Standalone comments refer to the next line
+            line += 1
+        self.comments[line] = node.comment

libcst/codemod/visitors/_gather_string_annotation_names.py

@@ -0,0 +1,65 @@
+# 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.
+
+from typing import Union, Set, cast
+
+import libcst as cst
+import libcst.matchers as m
+from libcst.codemod._context import CodemodContext
+from libcst.codemod._visitor import ContextAwareVisitor
+from libcst.metadata import QualifiedNameProvider
+
+FUNCS_CONSIDERED_AS_STRING_ANNOTATIONS = {"typing.TypeVar"}
+ANNOTATION_MATCHER = m.Annotation() | m.Call(
+    metadata=m.MatchMetadataIfTrue(
+        QualifiedNameProvider,
+        lambda qualnames: any(
+            qn.name in FUNCS_CONSIDERED_AS_STRING_ANNOTATIONS for qn in qualnames
+        ),
+    )
+)
+
+
+class GatherNamesFromStringAnnotationsVisitor(ContextAwareVisitor):
+    METADATA_DEPENDENCIES = (QualifiedNameProvider,)
+
+    def __init__(self, context: CodemodContext) -> None:
+        super().__init__(context)
+
+        self.names: Set[str] = set()
+
+    @m.call_if_inside(ANNOTATION_MATCHER)
+    @m.visit(m.ConcatenatedString())
+    def handle_any_string(
+        self, node: Union[cst.SimpleString, cst.ConcatenatedString]
+    ) -> None:
+        value = node.evaluated_value
+        if value is None:
+            return
+        mod = cst.parse_module(value)
+        extracted_nodes = m.extractall(
+            mod,
+            m.Name(value=m.SaveMatchedNode(m.DoNotCare(), "name"))
+            | m.SaveMatchedNode(m.Attribute(), "attribute"),
+        )
+        # This captures a bit more than necessary. For attributes, we capture the inner
+        # Name twice.
+        names = {
+            cast(str, values["name"]) for values in extracted_nodes if "name" in values
+        } | {
+            name
+            for values in extracted_nodes
+            if "attribute" in values
+            for name, _ in cst.metadata.scope_provider._gen_dotted_names(
+                cast(cst.Attribute, values["attribute"])
+            )
+        }
+        self.names.update(names)
+
+    @m.call_if_inside(ANNOTATION_MATCHER)
+    @m.call_if_not_inside(m.ConcatenatedString())
+    @m.visit(m.SimpleString())
+    def handle_simple_string(self, node: cst.SimpleString) -> None:
+        self.handle_any_string(node)

libcst/codemod/visitors/_gather_unused_imports.py

@@ -0,0 +1,121 @@
+# 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.
+#
+
+from typing import Dict, List, Optional, Sequence, Set, Tuple, Union, Iterable
+
+import libcst as cst
+import libcst.matchers as m
+from libcst.codemod._context import CodemodContext
+from libcst.codemod._visitor import ContextAwareVisitor
+from libcst.metadata import QualifiedNameProvider, ScopeProvider, PositionProvider
+from libcst.codemod.visitors._gather_exports import GatherExportsVisitor
+from libcst.codemod.visitors._gather_comments import GatherCommentsVisitor
+from libcst.codemod.visitors._gather_string_annotation_names import (
+    GatherNamesFromStringAnnotationsVisitor,
+)
+from libcst.metadata.scope_provider import _gen_dotted_names
+
+
+DEFAULT_SUPPRESS_COMMENT_REGEX = (
+    r".*\W(lint-ignore: ?unused-import|noqa|lint-ignore: ?F401)(\W.*)?$"
+)
+
+
+class GatherUnusedImportsVisitor(ContextAwareVisitor):
+
+    SUPPRESS_COMMENT_REGEX_CONTEXT_KEY = f"GatherUnusedImportsVisitor.suppress_regex"
+    METADATA_DEPENDENCIES = (
+        *GatherCommentsVisitor.METADATA_DEPENDENCIES,
+        *GatherNamesFromStringAnnotationsVisitor.METADATA_DEPENDENCIES,
+        PositionProvider,
+        ScopeProvider,
+    )
+
+    def __init__(self, context: CodemodContext) -> None:
+        super().__init__(context)
+
+        self._string_annotation_names: Set[str] = set()
+        self._ignored_lines: Set[int] = set()
+        self._exported_names: Set[str] = set()
+        self.unused_imports: Set[
+            Tuple[cst.ImportAlias, Union[cst.Import, cst.ImportFrom]]
+        ] = set()
+
+    def visit_Module(self, node: cst.Module) -> bool:
+        export_collector = GatherExportsVisitor(self.context)
+        node.visit(export_collector)
+        self._exported_names = export_collector.explicit_exported_objects
+        comment_visitor = GatherCommentsVisitor(
+            self.context,
+            self.context.scratch.get(
+                self.SUPPRESS_COMMENT_REGEX_CONTEXT_KEY, DEFAULT_SUPPRESS_COMMENT_REGEX,
+            ),
+        )
+        node.visit(comment_visitor)
+        self._ignored_lines = set(comment_visitor.comments.keys())
+        annotation_visitor = GatherNamesFromStringAnnotationsVisitor(self.context)
+        node.visit(annotation_visitor)
+        self._string_annotation_names = annotation_visitor.names
+        return True
+
+    @m.visit(
+        m.Import()
+        | m.ImportFrom(
+            module=m.DoesNotMatch(m.Name("__future__")),
+            names=m.DoesNotMatch(m.ImportStar()),
+        )
+    )
+    def handle_import(self, node: Union[cst.Import, cst.ImportFrom]) -> None:
+        assert not isinstance(node.names, cst.ImportStar)  # hello, type checker
+
+        node_start = self.get_metadata(PositionProvider, node).start.line
+        if node_start in self._ignored_lines:
+            return
+
+        for alias in node.names:
+            position = self.get_metadata(PositionProvider, alias)
+            lines = set(range(position.start.line, position.end.line + 1))
+            if lines.isdisjoint(self._ignored_lines):
+                self.unused_imports.add((alias, node))
+
+    def leave_Module(self, original_node: cst.Module) -> None:
+        self.unused_imports = self.filter_unused_imports(self.unused_imports)
+
+    def filter_unused_imports(
+        self,
+        candidates: Iterable[Tuple[cst.ImportAlias, Union[cst.Import, cst.ImportFrom]]],
+    ) -> Set[Tuple[cst.ImportAlias, Union[cst.Import, cst.ImportFrom]]]:
+        unused_imports = set()
+        for (alias, parent) in candidates:
+            scope = self.get_metadata(ScopeProvider, parent)
+            if scope is None:
+                continue
+            if not self.is_in_use(scope, alias):
+                unused_imports.add((alias, parent))
+        return unused_imports
+
+    def is_in_use(self, scope: cst.metadata.Scope, alias: cst.ImportAlias) -> bool:
+        asname = alias.asname
+        names = _gen_dotted_names(
+            cst.ensure_type(asname.name, cst.Name) if asname is not None else alias.name
+        )
+
+        for name_or_alias, _ in names:
+            if (
+                name_or_alias in self._exported_names
+                or name_or_alias in self._string_annotation_names
+            ):
+                return True
+
+            for assignment in scope[name_or_alias]:
+                if (
+                    isinstance(assignment, cst.metadata.Assignment)
+                    and isinstance(assignment.node, (cst.ImportFrom, cst.Import))
+                    and len(assignment.references) > 0
+                ):
+                    return True
+        return False
+

libcst/codemod/visitors/_remove_imports.py

@@ -8,7 +8,7 @@
 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.codemod.visitors._gather_unused_imports import GatherUnusedImportsVisitor
 from libcst.helpers import get_absolute_module_for_import, get_full_name_for_node
 from libcst.metadata import Assignment, Scope, ScopeProvider
 from libcst.metadata.scope_provider import _gen_dotted_names
@@ -173,7 +173,7 @@ def leave_AnnAssign(
     """
 
     CONTEXT_KEY = "RemoveImportsVisitor"
-    METADATA_DEPENDENCIES = (ScopeProvider,)
+    METADATA_DEPENDENCIES = (*GatherUnusedImportsVisitor.METADATA_DEPENDENCIES,)
 
     @staticmethod
     def _get_imports_from_context(
@@ -279,48 +279,24 @@ def __init__(
             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))
+        self._unused_imports: Dict[
+            cst.ImportAlias, Union[cst.Import, cst.ImportFrom]
+        ] = {}
 
     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
-        names = _gen_dotted_names(
-            cst.ensure_type(asname.name, cst.Name) if asname is not None else alias.name
-        )
-
-        for name_or_alias, _ in names:
-            if name_or_alias in self.exported_objects:
-                return True
-
-            for assignment in scope[name_or_alias]:
-                if (
-                    isinstance(assignment, Assignment)
-                    and isinstance(assignment.node, (cst.ImportFrom, cst.Import))
-                    and len(assignment.references) > 0
-                ):
-                    return True
-        return False
+        visitor = GatherUnusedImportsVisitor(self.context)
+        node.visit(visitor)
+        self._unused_imports = {k: v for (k, v) in visitor.unused_imports}
 
     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:
@@ -339,7 +315,7 @@ def leave_Import(
 
             # 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):
+            if import_alias not in self._unused_imports:
                 names_to_keep.append(import_alias)
                 continue
 
@@ -363,13 +339,6 @@ def leave_Import(
     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.
@@ -400,7 +369,7 @@ def leave_ImportFrom(
 
             # 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):
+            if import_alias not in self._unused_imports:
                 names_to_keep.append(import_alias)
                 continue