Infer: Allow combining two ModuleAnnotations

Summary:
The `__add__` function combines the annotations in two
`ModuleAnnotations` instances, preferring annotations from
the left when there are collisions.

The current purpose is to be able to combine together
existing annotations with the results of pyre infer (preferring infer
output when both are available) when making full stub files.

Reviewed By: pradeep90

Differential Revision: D28935186

fbshipit-source-id: c960ecd4f7976496a3fd96d13aea4aa2384816f1

Author: stroxler

client/commands/infer_v2.py

@@ -372,6 +372,45 @@ def from_module(path: str, module: libcst.Module) -> ModuleAnnotations:
             methods=collector.methods,
         )
 
+    def __add__(self, other: ModuleAnnotations) -> ModuleAnnotations:
+        """
+        Combine two sets of annotations for the same module, preferring
+        those from the right where there are collisions.
+
+        One use for this is to add existing annotations to those from infer
+        when generating full stub files.
+        """
+        if self.path != other.path:
+            raise ValueError(
+                "Cannot add ModuleAnnotations from different paths "
+                + f"{self.path!r} vs {other.path!r}"
+            )
+        other_names: set[str] = {
+            annotation.name
+            for annotation in
+            # pyre-ignore[58] : list[A] + list[B] = list[A | B]
+            other.globals_ + other.attributes + other.functions + other.methods
+        }
+        return ModuleAnnotations(
+            path=self.path,
+            globals_=other.globals_
+            + [global_ for global_ in self.globals_ if global_.name not in other_names],
+            attributes=other.attributes
+            + [
+                attribute
+                for attribute in self.attributes
+                if attribute.name not in other_names
+            ],
+            functions=other.functions
+            + [
+                function
+                for function in self.functions
+                if function.name not in other_names
+            ],
+            methods=other.methods
+            + [method for method in self.methods if method.name not in other_names],
+        )
+
     def filter_for_complete(self) -> ModuleAnnotations:
         return ModuleAnnotations(
             path=self.path,

client/commands/tests/infer_v2_test.py

@@ -569,10 +569,15 @@ 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:
+    def _assert_stubs(self, code: str, expected: str, data: dict | None = None) -> None:
         module_annotations = ModuleAnnotations.from_module(
             path=PATH, module=libcst.parse_module(textwrap.dedent(code))
         )
+        if data is not None:
+            module_annotations += _create_test_module_annotations(
+                data=data,
+                complete_only=False,
+            )
         actual = module_annotations.to_stubs()
         _assert_stubs_equal(actual, expected)
 
@@ -642,3 +647,68 @@ class Foo:
                 x: Any = ...
             """,
         )
+
+    def test_stubs_combining_annotations(self) -> None:
+        self._assert_stubs(
+            """
+            x: object = 1 + 1
+            y: int = 1 + 1
+            def f() -> object: return 10
+            def g() -> str: return "hello"
+            class Foo:
+                x: object = 10
+                y: int = 10
+                def f(self) -> object: return 10
+                def g(self) -> str: return "hello"
+            """,
+            data={
+                "globals": [
+                    {
+                        "name": "test.x",
+                        "annotation": "int",
+                    }
+                ],
+                "attributes": [
+                    {
+                        "name": "test.Foo.x",
+                        "parent": "test.Foo",
+                        "annotation": "int",
+                    }
+                ],
+                "defines": [
+                    {
+                        "name": "test.f",
+                        "parent": None,
+                        "return": "int",
+                        "parameters": [],
+                        "decorators": [],
+                        "async": False,
+                    },
+                    {
+                        "name": "test.Foo.f",
+                        "parent": "Foo",
+                        "return": "int",
+                        "parameters": [
+                            {
+                                "name": "self",
+                                "annotation": None,
+                                "value": None,
+                            }
+                        ],
+                        "decorators": [],
+                        "async": False,
+                    },
+                ],
+            },
+            expected="""\
+            x: int = ...
+            y: int = ...
+            def f() -> int: ...
+            def g() -> str: ...
+            class Foo:
+                x: int = ...
+                y: int = ...
+                def f(self) -> int: ...
+                def g(self) -> str: ...
+            """,
+        )