Repository: quora/pyanalyze
Pull request #324 (2021-12-08 15:59:00)

overload resolution fixes


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index d7da1d1..60d398e 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,7 +2,7 @@
 
 ## Unreleased
 
-- Improved overload matching algorithm (#321)
+- Improved overload matching algorithm (#321) (#324)
 - Support runtime overloaded functions with `pyanalyze.extensions.overload` (#318)
 - Internal support for overloaded functions (#316)
 - Support `TypeVar` bounds and constraints (#315)
diff --git a/pyanalyze/signature.py b/pyanalyze/signature.py
index ce00050..a116928 100644
--- a/pyanalyze/signature.py
+++ b/pyanalyze/signature.py
@@ -1573,6 +1573,9 @@ class OverloadedSignature:
         then used to construct a new :class:`ActualArguments` object, which ends up in
         :attr:`ImplReturn.remaining_arguments`.
 
+        An overload that matches without requiring use of ``Any`` or
+        union decomposition is called a "clean match".
+
         """
         actual_args = preprocess_args(args, visitor, node)
         if actual_args is None:
@@ -1581,24 +1584,45 @@ class OverloadedSignature:
         errors_per_overload = []
         any_rets: List[ImplReturn] = []
         union_rets: List[ImplReturn] = []
-        for sig in self.signatures:
+        union_and_any_rets: List[ImplReturn] = []
+        last = len(self.signatures) - 1
+        for i, sig in enumerate(self.signatures):
             with visitor.catch_errors() as caught_errors:
                 ret = sig.check_call_preprocessed(
-                    actual_args, visitor, node, is_overload=True
+                    actual_args,
+                    visitor,
+                    node,
+                    # We set is_overload to False for the last overload
+                    # because we can't do union decomposition on the last one:
+                    # there's no other overload that could handle the remaining
+                    # union members.
+                    is_overload=i != last,
                 )
+            errors_per_overload.append(caught_errors)
             if ret.is_error:
-                errors_per_overload.append(caught_errors)
-            elif ret.used_any_for_match:
-                any_rets.append(ret)
+                continue
             elif ret.remaining_arguments is not None:
-                union_rets.append(ret)
+                if ret.used_any_for_match:
+                    # If an overload used both Any and union decomposition,
+                    # we treat it differently from either: unlike Any matches
+                    # it's not enough to prevent an error later (because it's
+                    # not a full match), but unlike union matches it's enough
+                    # to trigger an Any return type.
+                    union_and_any_rets.append(ret)
+                else:
+                    union_rets.append(ret)
                 actual_args = ret.remaining_arguments
+            elif ret.used_any_for_match:
+                any_rets.append(ret)
             else:
                 # We got a clean match!
-                return self._unite_rets(any_rets, union_rets, ret)
+                return self._unite_rets(any_rets, union_and_any_rets, union_rets, ret)
 
-        if any_rets or union_rets:
-            return self._unite_rets(any_rets, union_rets)
+        if any_rets:
+            # We don't do this if we have union_rets, because if we got here, we
+            # didn't get any clean matches. Therefore, we must have some remaining
+            # union members we haven't handled.
+            return self._unite_rets(any_rets, union_and_any_rets, union_rets)
 
         # None of the signatures matched
         errors = list(itertools.chain.from_iterable(errors_per_overload))
@@ -1616,11 +1640,17 @@ class OverloadedSignature:
     def _unite_rets(
         self,
         any_rets: Sequence[ImplReturn],
+        union_and_any_rets: Sequence[ImplReturn],
         union_rets: Sequence[ImplReturn],
         clean_ret: Optional[ImplReturn] = None,
     ) -> ImplReturn:
-        if any_rets:
-            if len(any_rets) == 1 and not union_rets and clean_ret is None:
+        if any_rets or union_and_any_rets:
+            if (
+                len(any_rets) == 1
+                and not union_rets
+                and not union_and_any_rets
+                and clean_ret is None
+            ):
                 return any_rets[0]
             return ImplReturn(AnyValue(AnySource.multiple_overload_matches))
         elif union_rets:


Test patch:
diff --git a/pyanalyze/test_signature.py b/pyanalyze/test_signature.py
index ee6ab64..0562fd0 100644
--- a/pyanalyze/test_signature.py
+++ b/pyanalyze/test_signature.py
@@ -1066,3 +1066,43 @@ class TestOverload(TestNameCheckVisitorBase):
             assert_is_value(val5, KnownValue(3) | KnownValue(2))
             val6 = overloaded(explicit)  # pyright: Literal[2], mypy: Any
             assert_is_value(val6, AnyValue(AnySource.multiple_overload_matches))
+
+    @assert_passes()
+    def test_any_and_union(self):
+        from pyanalyze.extensions import overload
+        from typing import List, Any, Union
+        from typing_extensions import Literal
+
+        @overload
+        def overloaded1(x: Any, y: str) -> Literal[2]:
+            pass
+
+        @overload
+        def overloaded1(x: str, y: int) -> Literal[3]:
+            pass
+
+        def overloaded1(x: object, y: object) -> Literal[2, 3]:
+            raise NotImplementedError
+
+        @overload
+        def overloaded2(x: object, y: str) -> Literal[2]:
+            pass
+
+        @overload
+        def overloaded2(x: str, y: int) -> Literal[3]:
+            pass
+
+        def overloaded2(x: object, y: object) -> Literal[2, 3]:
+            raise NotImplementedError
+
+        def capybara(
+            unannotated,
+            int_or_str: Union[int, str],
+            int_or_str_or_float: Union[int, str, float],
+        ):
+            val1 = overloaded1(unannotated, int_or_str)
+            assert_is_value(val1, AnyValue(AnySource.multiple_overload_matches))
+            val2 = overloaded2(unannotated, int_or_str)
+            assert_is_value(val2, AnyValue(AnySource.multiple_overload_matches))
+            val3 = overloaded2("x", int_or_str_or_float)  # E: incompatible_argument
+            assert_is_value(val3, AnyValue(AnySource.error))