Implement pyre infer (7/n) -- raw infer output to module annotations

Reviewed By: dkgi

Differential Revision: D30169045

fbshipit-source-id: 47fe64d8c6b036300241817e25455a2d1c633ceb

Author: grievejia

client/commands/v2/infer.py

@@ -211,6 +211,154 @@ class RawInferOutputParsingError(Exception):
     pass
 
 
+@dataclasses.dataclass(frozen=True)
+class StubGenerationOptions:
+    annotate_attributes: bool
+    use_future_annotations: bool
+    dequalify: bool
+
+
+@dataclasses.dataclass(frozen=True)
+class TypeAnnotation:
+    annotation: Optional[str]
+    dequalify: bool
+
+    @staticmethod
+    def from_raw(
+        annotation: Optional[str], options: StubGenerationOptions
+    ) -> "TypeAnnotation":
+        return TypeAnnotation(annotation=annotation, dequalify=options.dequalify)
+
+    @property
+    def missing(self) -> bool:
+        return self.annotation is None
+
+
+@dataclasses.dataclass(frozen=True)
+class Parameter:
+    name: str
+    annotation: TypeAnnotation
+    value: Optional[str]
+
+
+@dataclasses.dataclass(frozen=True)
+class FunctionAnnotation:
+    name: str
+    return_annotation: TypeAnnotation
+    parameters: Sequence[Parameter]
+    decorators: Sequence[str]
+    is_async: bool
+
+
+@dataclasses.dataclass(frozen=True)
+class MethodAnnotation(FunctionAnnotation):
+    parent: str
+
+
+@dataclasses.dataclass(frozen=True)
+class FieldAnnotation:
+    name: str
+    annotation: TypeAnnotation
+
+
+@dataclasses.dataclass(frozen=True)
+class GlobalAnnotation(FieldAnnotation):
+    pass
+
+
+@dataclasses.dataclass(frozen=True)
+class AttributeAnnotation(FieldAnnotation):
+    parent: str
+
+
+@dataclasses.dataclass(frozen=True)
+class ModuleAnnotations:
+    path: str
+    options: StubGenerationOptions
+    globals_: List[GlobalAnnotation] = dataclasses.field(default_factory=list)
+    attributes: List[AttributeAnnotation] = dataclasses.field(default_factory=list)
+    functions: List[FunctionAnnotation] = dataclasses.field(default_factory=list)
+    methods: List[MethodAnnotation] = dataclasses.field(default_factory=list)
+
+    @staticmethod
+    def from_infer_output(
+        path: str,
+        infer_output: RawInferOutput,
+        options: StubGenerationOptions,
+    ) -> "ModuleAnnotations":
+        def type_annotation(annotation: Optional[str]) -> TypeAnnotation:
+            return TypeAnnotation.from_raw(
+                annotation,
+                options=options,
+            )
+
+        return ModuleAnnotations(
+            path=path,
+            globals_=[
+                GlobalAnnotation(
+                    name=global_.name, annotation=type_annotation(global_.annotation)
+                )
+                for global_ in infer_output.global_annotations
+            ],
+            attributes=[
+                AttributeAnnotation(
+                    parent=attribute.parent,
+                    name=attribute.name,
+                    annotation=type_annotation(attribute.annotation),
+                )
+                for attribute in infer_output.attribute_annotations
+            ]
+            if options.annotate_attributes
+            else [],
+            functions=[
+                FunctionAnnotation(
+                    name=define.name,
+                    return_annotation=type_annotation(define.return_),
+                    parameters=[
+                        Parameter(
+                            name=parameter.name,
+                            annotation=type_annotation(parameter.annotation),
+                            value=parameter.value,
+                        )
+                        for parameter in define.parameters
+                    ],
+                    decorators=define.decorators,
+                    is_async=define.is_async,
+                )
+                for define in infer_output.define_annotations
+                if define.parent is None
+            ],
+            methods=[
+                MethodAnnotation(
+                    parent=define.parent,
+                    name=define.name,
+                    return_annotation=type_annotation(define.return_),
+                    parameters=[
+                        Parameter(
+                            name=parameter.name,
+                            annotation=type_annotation(parameter.annotation),
+                            value=parameter.value,
+                        )
+                        for parameter in define.parameters
+                    ],
+                    decorators=define.decorators,
+                    is_async=define.is_async,
+                )
+                for define in infer_output.define_annotations
+                if define.parent is not None
+            ],
+            options=options,
+        )
+
+    def is_empty(self) -> bool:
+        return (
+            len(self.globals_)
+            + len(self.attributes)
+            + len(self.functions)
+            + len(self.methods)
+        ) == 0
+
+
 def create_infer_arguments(
     configuration: configuration_module.Configuration,
     infer_arguments: command_arguments.InferArguments,

client/commands/v2/tests/infer_test.py

@@ -24,6 +24,13 @@
     RawDefineAnnotation,
     RawInferOutput,
     RawInferOutputParsingError,
+    StubGenerationOptions,
+    TypeAnnotation,
+    FunctionAnnotation,
+    MethodAnnotation,
+    GlobalAnnotation,
+    AttributeAnnotation,
+    ModuleAnnotations,
 )
 
 
@@ -332,3 +339,163 @@ def assert_not_parsed(input: str) -> None:
                 ]
             ),
         )
+
+
+class ModuleAnnotationTest(testslide.TestCase):
+    def test_module_annotations_from_infer_output(self) -> None:
+        def assert_result(
+            path: str,
+            infer_output: RawInferOutput,
+            options: StubGenerationOptions,
+            expected: ModuleAnnotations,
+        ) -> None:
+            self.assertEqual(
+                ModuleAnnotations.from_infer_output(path, infer_output, options),
+                expected,
+            )
+
+        default_path = "test.py"
+        default_options = StubGenerationOptions(
+            annotate_attributes=False, use_future_annotations=False, dequalify=False
+        )
+
+        assert_result(
+            path=default_path,
+            infer_output=RawInferOutput(),
+            options=default_options,
+            expected=ModuleAnnotations(path=default_path, options=default_options),
+        )
+
+        assert_result(
+            path=default_path,
+            infer_output=RawInferOutput(
+                define_annotations=[
+                    RawDefineAnnotation(
+                        name="test.Foo.foo",
+                        parent="test.Foo",
+                        location=RawAnnotationLocation(
+                            qualifier="test", path="test.py", line=1
+                        ),
+                        return_="int",
+                        decorators=["derp"],
+                        is_async=True,
+                    ),
+                    RawDefineAnnotation(
+                        name="test.bar",
+                        location=RawAnnotationLocation(
+                            qualifier="test", path="test.py", line=2
+                        ),
+                    ),
+                ],
+            ),
+            options=default_options,
+            expected=ModuleAnnotations(
+                path=default_path,
+                options=default_options,
+                functions=[
+                    FunctionAnnotation(
+                        name="test.bar",
+                        return_annotation=TypeAnnotation.from_raw(
+                            None, options=default_options
+                        ),
+                        parameters=[],
+                        decorators=[],
+                        is_async=False,
+                    )
+                ],
+                methods=[
+                    MethodAnnotation(
+                        parent="test.Foo",
+                        name="test.Foo.foo",
+                        return_annotation=TypeAnnotation.from_raw(
+                            "int", options=default_options
+                        ),
+                        parameters=[],
+                        decorators=["derp"],
+                        is_async=True,
+                    )
+                ],
+            ),
+        )
+
+        assert_result(
+            path=default_path,
+            infer_output=RawInferOutput(
+                global_annotations=[
+                    RawGlobalAnnotation(
+                        name="x",
+                        location=RawAnnotationLocation(
+                            qualifier="test", path="test.py", line=3
+                        ),
+                        annotation="int",
+                    )
+                ],
+            ),
+            options=default_options,
+            expected=ModuleAnnotations(
+                path=default_path,
+                options=default_options,
+                globals_=[
+                    GlobalAnnotation(
+                        name="x",
+                        annotation=TypeAnnotation.from_raw(
+                            "int", options=default_options
+                        ),
+                    )
+                ],
+            ),
+        )
+
+        assert_result(
+            path=default_path,
+            infer_output=RawInferOutput(
+                attribute_annotations=[
+                    RawAttributeAnnotation(
+                        parent="Foo",
+                        name="x",
+                        location=RawAnnotationLocation(
+                            qualifier="test", path="test.py", line=3
+                        ),
+                        annotation="int",
+                    )
+                ],
+            ),
+            options=default_options,
+            expected=ModuleAnnotations(
+                path=default_path,
+                options=default_options,
+            ),
+        )
+
+        annotate_attribute_options = StubGenerationOptions(
+            annotate_attributes=True, use_future_annotations=False, dequalify=False
+        )
+        assert_result(
+            path=default_path,
+            infer_output=RawInferOutput(
+                attribute_annotations=[
+                    RawAttributeAnnotation(
+                        parent="Foo",
+                        name="x",
+                        location=RawAnnotationLocation(
+                            qualifier="test", path="test.py", line=3
+                        ),
+                        annotation="int",
+                    )
+                ],
+            ),
+            options=annotate_attribute_options,
+            expected=ModuleAnnotations(
+                path=default_path,
+                options=annotate_attribute_options,
+                attributes=[
+                    AttributeAnnotation(
+                        parent="Foo",
+                        name="x",
+                        annotation=TypeAnnotation.from_raw(
+                            "int", options=annotate_attribute_options
+                        ),
+                    )
+                ],
+            ),
+        )