Infer: Add support for existing annotations to infer_v2

Summary:
The existing infer has the ability to generate complete stub files,
which include all explicitly annotated attributes and globals
as well as all functions and methods with explicit return annotations.

This commit adds support for extracting existing annotations from a
module using `libcst`; the next diff adds the ability to combine them
with annotations from infer.

Reviewed By: pradeep90

Differential Revision: D28918138

fbshipit-source-id: d38be3550e68a1ec31e1c57453979d68163c469d

Author: stroxler

client/commands/infer_v2.py

@@ -14,12 +14,12 @@
 from pathlib import Path
 from typing import cast, Dict, List, Optional, Sequence
 
+import libcst
 from typing_extensions import Final
 from typing_extensions import TypeAlias
 
 from .. import command_arguments, log
 from ..analysis_directory import AnalysisDirectory, resolve_analysis_directory
-from ..annotation_collector import AnnotationCollector
 from ..configuration import Configuration
 from .check import Check
 from .infer import dequalify_and_fix_pathlike, split_imports
@@ -176,6 +176,114 @@ class AttributeAnnotation(FieldAnnotation):
     parent: str
 
 
+class AnnotationCollector(libcst.CSTVisitor):
+    def __init__(self, file_path: str) -> None:
+        self.qualifier: str = file_path.replace(".py", "")
+        self.globals_: list[GlobalAnnotation] = []
+        self.attributes: list[AttributeAnnotation] = []
+        self.functions: list[FunctionAnnotation] = []
+        self.methods: list[MethodAnnotation] = []
+        self.class_name: list[str] = []
+
+    def visit_ClassDef(self, node: libcst.ClassDef) -> None:
+        self.class_name.append(node.name.value)
+
+    def leave_ClassDef(self, original_node: libcst.ClassDef) -> None:
+        self.class_name.remove(original_node.name.value)
+
+    def visit_AnnAssign(self, node: libcst.AnnAssign) -> None:
+        target_name = self._code_for_node(node.target)
+        if target_name is None:
+            return
+        parent = ".".join(self.class_name) if self.class_name else None
+        name = self._qualified_name(
+            parent=parent,
+            target_name=target_name,
+        )
+        annotation = TypeAnnotation(self._code_for_node(node.annotation.annotation))
+        if parent is None:
+            self.globals_.append(
+                GlobalAnnotation(
+                    name=name,
+                    annotation=annotation,
+                )
+            )
+        else:
+            self.attributes.append(
+                AttributeAnnotation(
+                    parent=parent,
+                    name=name,
+                    annotation=annotation,
+                )
+            )
+
+    def visit_FunctionDef(self, node: libcst.FunctionDef) -> bool:
+        if node.returns is None:
+            return False
+        parent = ".".join(self.class_name) if self.class_name else None
+        name = self._qualified_name(
+            parent=parent,
+            target_name=node.name.value,
+        )
+        return_annotation = TypeAnnotation(self._code_for_node(node.returns.annotation))
+        parameters = [
+            Parameter(
+                name=parameter.name.value,
+                annotation=TypeAnnotation(
+                    self._code_for_node(
+                        parameter.annotation.annotation
+                        if parameter.annotation
+                        else None
+                    )
+                ),
+                value=self._code_for_node(parameter.default),
+            )
+            for parameter in node.params.params
+        ]
+        decorators = [
+            decorator
+            for decorator in (
+                self._code_for_node(decorator) for decorator in node.decorators
+            )
+            if decorator is not None
+        ]
+        is_async = node.asynchronous is not None
+        if parent is None:
+            self.functions.append(
+                FunctionAnnotation(
+                    name=name,
+                    return_annotation=return_annotation,
+                    parameters=parameters,
+                    decorators=decorators,
+                    is_async=is_async,
+                )
+            )
+        else:
+            self.methods.append(
+                MethodAnnotation(
+                    parent=parent,
+                    name=name,
+                    return_annotation=return_annotation,
+                    parameters=parameters,
+                    decorators=decorators,
+                    is_async=is_async,
+                )
+            )
+        return False
+
+    # utility methods
+
+    def _qualified_name(self, parent: str | None, target_name: str) -> str:
+        prefix = f"{self.qualifier}."
+        if parent is not None:
+            prefix += f"{parent}."
+        return prefix + target_name
+
+    @staticmethod
+    def _code_for_node(node: Optional[libcst.CSTNode]) -> Optional[str]:
+        return None if node is None else libcst.parse_module("").code_for_node(node)
+
+
 @dataclass
 class ModuleAnnotations:
     path: str
@@ -252,6 +360,18 @@ def from_infer_output(
             ],
         )
 
+    @staticmethod
+    def from_module(path: str, module: libcst.Module) -> ModuleAnnotations:
+        collector = AnnotationCollector(file_path=path)
+        module.visit(collector)
+        return ModuleAnnotations(
+            path=path,
+            globals_=collector.globals_,
+            attributes=collector.attributes,
+            functions=collector.functions,
+            methods=collector.methods,
+        )
+
     def filter_for_complete(self) -> ModuleAnnotations:
         return ModuleAnnotations(
             path=self.path,

client/commands/tests/infer_v2_test.py

@@ -11,6 +11,8 @@
 from pathlib import Path
 from typing import cast
 
+import libcst
+
 from ...commands.infer_v2 import (
     _create_module_annotations,
     RawInferOutput,
@@ -564,3 +566,79 @@ def test_stubs_no_typing_import(self) -> None:
             def with_params(y=7, x: List[int] = [5]) -> Union[int, str]: ...
             """,
         )
+
+
+class ExistingAnnotationsTest(unittest.TestCase):
+    def _assert_stubs(self, code: str, expected: str) -> None:
+        module_annotations = ModuleAnnotations.from_module(
+            path=PATH, module=libcst.parse_module(textwrap.dedent(code))
+        )
+        actual = module_annotations.to_stubs()
+        _assert_stubs_equal(actual, expected)
+
+    def test_stubs_from_existing_annotations(self) -> None:
+        self._assert_stubs(
+            """
+            def foo() -> int:
+                return 1 + 1
+            """,
+            "def foo() -> int: ...",
+        )
+
+        # methods
+        self._assert_stubs(
+            """
+            class Foo:
+                def bar(self, x: int) -> Union[int, str]:
+                    return ""
+            """,
+            """\
+            class Foo:
+                def bar(self, x: int) -> Union[int, str]: ...
+            """,
+        )
+
+        # with async
+        self._assert_stubs(
+            """
+            async def foo() -> int:
+                return 1 + 1
+            """,
+            "async def foo() -> int: ...",
+        )
+
+        # with decorators
+        self._assert_stubs(
+            """
+            @click
+            def foo() -> int:
+                return 1 + 1
+            """,
+            "@@click\n\ndef foo() -> int: ...",
+        )
+
+        # globals
+        self._assert_stubs(
+            """
+            x: int = 10
+            """,
+            "x: int = ...",
+        )
+
+        # attributes
+        self._assert_stubs(
+            # TODO (T92336996)
+            # libcst does not produce fully-qualified typenames when extracting
+            # annotations (unlike the pyre parser). As a result, we're producing
+            # incorrect stubs here.
+            """
+            from typing import Any
+
+            class Foo:
+                x: Any = 10
+            """,
+            """\
+            class Foo:
+                x: Any = ...
+            """,
+        )