feat: Support merging overload annotations into implementation

Issue-442: #442
PR-443: #443

Author: pawamoy

packages/griffelib/src/griffe/_internal/expressions.py

@@ -251,6 +251,22 @@ def is_generator(self) -> bool:
         """Whether this expression is a generator."""
         return isinstance(self, ExprSubscript) and self.canonical_name == "Generator"
 
+    @staticmethod
+    def _to_binop(elements: Sequence[Expr], op: str) -> ExprBinOp:
+        if len(elements) == 2:  # noqa: PLR2004
+            left, right = elements
+            if isinstance(left, Expr):
+                left = left.modernize()
+            if isinstance(right, Expr):
+                right = right.modernize()
+            return ExprBinOp(left=left, operator=op, right=right)
+
+        left = ExprSubscript._to_binop(elements[:-1], op=op)
+        right = elements[-1]
+        if isinstance(right, Expr):
+            right = right.modernize()
+        return ExprBinOp(left=left, operator=op, right=right)
+
 
 @dataclass(eq=True, slots=True)
 class ExprAttribute(Expr):
@@ -888,22 +904,6 @@ def iterate(self, *, flat: bool = True) -> Iterator[str | Expr]:
         yield from _yield(self.slice, flat=flat, outer_precedence=_OperatorPrecedence.NONE)
         yield "]"
 
-    @staticmethod
-    def _to_binop(elements: Sequence[Expr], op: str) -> ExprBinOp:
-        if len(elements) == 2:  # noqa: PLR2004
-            left, right = elements
-            if isinstance(left, Expr):
-                left = left.modernize()
-            if isinstance(right, Expr):
-                right = right.modernize()
-            return ExprBinOp(left=left, operator=op, right=right)
-
-        left = ExprSubscript._to_binop(elements[:-1], op=op)
-        right = elements[-1]
-        if isinstance(right, Expr):
-            right = right.modernize()
-        return ExprBinOp(left=left, operator=op, right=right)
-
     def modernize(self) -> ExprBinOp | ExprSubscript:
         if self.canonical_path == "typing.Union":
             return self._to_binop(self.slice.elements, op="|")  # type: ignore[union-attr]

packages/griffelib/src/griffe/_internal/merger.py

@@ -6,9 +6,12 @@
 from typing import TYPE_CHECKING
 
 from griffe._internal.exceptions import AliasResolutionError, CyclicAliasError
+from griffe._internal.expressions import Expr
 from griffe._internal.logger import logger
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from griffe._internal.models import Attribute, Class, Function, Module, Object, TypeAlias
 
 
@@ -60,10 +63,43 @@ def _merge_stubs_overloads(obj: Module | Class, stubs: Module | Class) -> None:
     for function_name, overloads in list(stubs.overloads.items()):
         if overloads:
             with suppress(KeyError):
-                obj.get_member(function_name).overloads = overloads
+                _merge_overload_annotations(obj.get_member(function_name), overloads)
         del stubs.overloads[function_name]
 
 
+def _merge_annotations(annotations: Sequence[Expr]) -> Expr | None:
+    if len(annotations) == 1:
+        return annotations[0]
+    if annotations:
+        return Expr._to_binop(annotations, op="|")
+    return None
+
+
+def _merge_overload_annotations(function: Function, overloads: list[Function]) -> None:
+    function.overloads = overloads
+    for parameter in function.parameters:
+        if parameter.annotation is None:
+            seen = set()
+            annotations = []
+            for overload in overloads:
+                with suppress(KeyError):
+                    annotation = overload.parameters[parameter.name].annotation
+                    str_annotation = str(annotation)
+                    if isinstance(annotation, Expr) and str_annotation not in seen:
+                        annotations.append(annotation)
+                        seen.add(str_annotation)
+            parameter.annotation = _merge_annotations(annotations)
+    if function.returns is None:
+        seen = set()
+        return_annotations = []
+        for overload in overloads:
+            str_annotation = str(overload.returns)
+            if isinstance(overload.returns, Expr) and str_annotation not in seen:
+                return_annotations.append(overload.returns)
+                seen.add(str_annotation)
+        function.returns = _merge_annotations(return_annotations)
+
+
 def _merge_stubs_members(obj: Module | Class, stubs: Module | Class) -> None:
     # Merge imports to later know if objects coming from the stubs were imported.
     obj.imports.update(stubs.imports)

tests/test_merger.py

@@ -73,3 +73,25 @@ def test_merge_attribute_values() -> None:
         },
     ) as pkg:
         assert str(pkg["__all__"].value) == "['hello']"
+
+
+def test_merge_overload_annotations() -> None:
+    """Assert that overload annotations are merged correctly."""
+    with temporary_visited_package(
+        "package",
+        {
+            "mod.py": "def func(x): ...",
+            "mod.pyi": """
+                from typing import overload
+
+                @overload
+                def func(x: int) -> int: ...
+
+                @overload
+                def func(x: float) -> float: ...
+                """,
+        },
+    ) as pkg:
+        func = pkg["mod.func"]
+        assert str(func.parameters["x"].annotation) == "int | float"
+        assert str(func.returns) == "int | float"