Repository: quora/pyanalyze
Pull request #318 (2021-12-06 17:02:00)

support runtime @overload
Another step toward #182.


Patch:
diff --git a/pyanalyze/__init__.py b/pyanalyze/__init__.py
index 8bc67b5..e6fbc47 100644
--- a/pyanalyze/__init__.py
+++ b/pyanalyze/__init__.py
@@ -37,6 +37,7 @@ used(assert_is_value)
 used(dump_value)
 used(extensions.LiteralOnly)
 used(extensions.NoAny)
+used(extensions.overload)
 used(value.UNRESOLVED_VALUE)  # keeping it around for now just in case
 used(reexport)
 used(checker)
diff --git a/pyanalyze/arg_spec.py b/pyanalyze/arg_spec.py
index 3754800..8f0384d 100644
--- a/pyanalyze/arg_spec.py
+++ b/pyanalyze/arg_spec.py
@@ -4,16 +4,25 @@ Implementation of extended argument specifications used by test_scope.
 
 """
 
+from .extensions import get_overloads
 from .annotations import Context, type_from_runtime
 from .config import Config
 from .find_unused import used
 from . import implementation
-from .safe import is_newtype, safe_hasattr, safe_in, safe_issubclass, is_typing_name
+from .safe import (
+    all_of_type,
+    is_newtype,
+    safe_hasattr,
+    safe_in,
+    safe_issubclass,
+    is_typing_name,
+)
 from .stacked_scopes import Composite, uniq_chain
 from .signature import (
     ANY_SIGNATURE,
     Impl,
     MaybeSignature,
+    OverloadedSignature,
     PropertyArgSpec,
     make_bound_method,
     SigParameter,
@@ -293,10 +302,16 @@ class ArgSpecCache:
         self, obj: object, impl: Optional[Impl] = None, is_asynq: bool = False
     ) -> MaybeSignature:
         """Constructs the Signature for a Python object."""
-        return self._cached_get_argspec(obj, impl, is_asynq)
+        return self._cached_get_argspec(
+            obj, impl, is_asynq, in_overload_resolution=False
+        )
 
     def _cached_get_argspec(
-        self, obj: object, impl: Optional[Impl], is_asynq: bool
+        self,
+        obj: object,
+        impl: Optional[Impl],
+        is_asynq: bool,
+        in_overload_resolution: bool,
     ) -> MaybeSignature:
         try:
             if obj in self.known_argspecs:
@@ -306,7 +321,9 @@ class ArgSpecCache:
         else:
             hashable = True
 
-        extended = self._uncached_get_argspec(obj, impl, is_asynq)
+        extended = self._uncached_get_argspec(
+            obj, impl, is_asynq, in_overload_resolution
+        )
         if extended is None:
             return None
 
@@ -315,22 +332,44 @@ class ArgSpecCache:
         return extended
 
     def _uncached_get_argspec(
-        self, obj: Any, impl: Optional[Impl], is_asynq: bool
+        self,
+        obj: Any,
+        impl: Optional[Impl],
+        is_asynq: bool,
+        in_overload_resolution: bool,
     ) -> MaybeSignature:
+        if not in_overload_resolution:
+            fq_name = _get_fully_qualified_name(obj)
+            if fq_name is not None:
+                overloads = get_overloads(fq_name)
+                if overloads:
+                    sigs = [
+                        self._cached_get_argspec(
+                            overload, impl, is_asynq, in_overload_resolution=True
+                        )
+                        for overload in overloads
+                    ]
+                    if all_of_type(sigs, Signature):
+                        return OverloadedSignature(sigs)
+
         if isinstance(obj, tuple) or hasattr(obj, "__getattr__"):
             return None  # lost cause
 
         # Cythonized methods, e.g. fn.asynq
         if is_dot_asynq_function(obj):
             try:
-                return self._cached_get_argspec(obj.__self__, impl, is_asynq)
+                return self._cached_get_argspec(
+                    obj.__self__, impl, is_asynq, in_overload_resolution
+                )
             except TypeError:
                 # some cythonized methods have __self__ but it is not a function
                 pass
 
         # for bound methods, see if we have an argspec for the unbound method
         if inspect.ismethod(obj) and obj.__self__ is not None:
-            argspec = self._cached_get_argspec(obj.__func__, impl, is_asynq)
+            argspec = self._cached_get_argspec(
+                obj.__func__, impl, is_asynq, in_overload_resolution
+            )
             return make_bound_method(argspec, Composite(KnownValue(obj.__self__)))
 
         if hasattr(obj, "fn") or hasattr(obj, "original_fn"):
@@ -343,7 +382,9 @@ class ArgSpecCache:
                 # e.g. certain extension classes
                 pass
             else:
-                return self._cached_get_argspec(original_fn, impl, is_asynq)
+                return self._cached_get_argspec(
+                    original_fn, impl, is_asynq, in_overload_resolution
+                )
 
         argspec = self.ts_finder.get_argspec(obj)
         if argspec is not None:
@@ -367,7 +408,9 @@ class ArgSpecCache:
         if inspect.isfunction(obj):
             if hasattr(obj, "inner"):
                 # @qclient.task_queue.exec_after_request() puts the original function in .inner
-                return self._cached_get_argspec(obj.inner, impl, is_asynq)
+                return self._cached_get_argspec(
+                    obj.inner, impl, is_asynq, in_overload_resolution
+                )
 
             inspect_sig = self._safe_get_signature(obj)
             if inspect_sig is None:
@@ -383,7 +426,9 @@ class ArgSpecCache:
 
         # decorator binders
         if _is_qcore_decorator(obj):
-            argspec = self._cached_get_argspec(obj.decorator, impl, is_asynq)
+            argspec = self._cached_get_argspec(
+                obj.decorator, impl, is_asynq, in_overload_resolution
+            )
             # wrap if it's a bound method
             if obj.instance is not None and argspec is not None:
                 return make_bound_method(argspec, Composite(KnownValue(obj.instance)))
@@ -452,7 +497,9 @@ class ArgSpecCache:
                     return self._make_any_sig(obj)
                 if method == obj:
                     return self._make_any_sig(obj)
-                argspec = self._cached_get_argspec(method, impl, is_asynq)
+                argspec = self._cached_get_argspec(
+                    method, impl, is_asynq, in_overload_resolution
+                )
                 return make_bound_method(argspec, Composite(KnownValue(obj.__self__)))
             inspect_sig = self._safe_get_signature(obj)
             if inspect_sig is not None:
@@ -468,7 +515,9 @@ class ArgSpecCache:
         if isinstance(obj, property):
             # If we know the getter, inherit its return value.
             if obj.fget:
-                fget_argspec = self._cached_get_argspec(obj.fget, impl, is_asynq)
+                fget_argspec = self._cached_get_argspec(
+                    obj.fget, impl, is_asynq, in_overload_resolution
+                )
                 if fget_argspec is not None and fget_argspec.has_return_value():
                     return PropertyArgSpec(obj, return_value=fget_argspec.return_value)
             return PropertyArgSpec(obj)
@@ -603,3 +652,9 @@ def _is_qcore_decorator(obj: object) -> bool:
     except Exception:
         # black.Line has an is_decorator attribute but it is not a method
         return False
+
+
+def _get_fully_qualified_name(obj: object) -> Optional[str]:
+    if hasattr(obj, "__module__") and hasattr(obj, "__qualname__"):
+        return f"{obj.__module__}.{obj.__qualname__}"
+    return None
diff --git a/pyanalyze/asynq_checker.py b/pyanalyze/asynq_checker.py
index 7a3f446..1d721f1 100644
--- a/pyanalyze/asynq_checker.py
+++ b/pyanalyze/asynq_checker.py
@@ -34,6 +34,7 @@ class FunctionInfo:
     is_classmethod: bool  # has @classmethod
     is_staticmethod: bool  # has @staticmethod
     is_decorated_coroutine: bool  # has @asyncio.coroutine
+    is_overload: bool  # typing.overload or pyanalyze.extensions.overload
     # a list of pairs of (decorator function, applied decorator function). These are different
     # for decorators that take arguments, like @asynq(): the first element will be the asynq
     # function and the second will be the result of calling asynq().
diff --git a/pyanalyze/extensions.py b/pyanalyze/extensions.py
index 61c0d95..86a3f43 100644
--- a/pyanalyze/extensions.py
+++ b/pyanalyze/extensions.py
@@ -7,9 +7,22 @@ Several type system extensions are used with the ``Annotated`` type from
 be gracefully ignored by other type checkers.
 
 """
+from collections import defaultdict
 from dataclasses import dataclass, field
 import pyanalyze
-from typing import Any, Container, Iterable, Tuple, List, Union, TypeVar, TYPE_CHECKING
+from typing import (
+    Any,
+    Callable,
+    Container,
+    Dict,
+    Iterable,
+    Tuple,
+    List,
+    Union,
+    TypeVar,
+    TYPE_CHECKING,
+    overload as real_overload,
+)
 from typing_extensions import Literal, NoReturn
 
 if TYPE_CHECKING:
@@ -336,3 +349,29 @@ def reveal_type(value: object) -> None:
 
     """
     pass
+
+
+_overloads: Dict[str, List[Callable[..., Any]]] = defaultdict(list)
+
+
+def get_overloads(fully_qualified_name: str) -> List[Callable[..., Any]]:
+    """Return all defined runtime overloads for this fully qualified name."""
+    return _overloads[fully_qualified_name]
+
+
+if TYPE_CHECKING:
+    from typing import overload as overload
+
+else:
+
+    def overload(func: Callable[..., Any]) -> Callable[..., Any]:
+        """A version of `typing.overload` that is inspectable at runtime.
+
+        If this decorator is used for a function `some_module.some_function`, calling
+        :func:`pyanalyze.extensiions.get_overloads("some_module.some_function")` will
+        return all the runtime overloads.
+
+        """
+        key = f"{func.__module__}.{func.__qualname__}"
+        _overloads[key].append(func)
+        return real_overload(func)
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index ea2541f..8d79d49 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -31,6 +31,7 @@ import sys
 import tempfile
 import traceback
 import types
+import typing
 from typing import (
     ClassVar,
     ContextManager,
@@ -67,7 +68,7 @@ from .boolability import Boolability, get_boolability
 from .checker import Checker
 from .config import Config
 from .error_code import ErrorCode, DISABLED_BY_DEFAULT, ERROR_DESCRIPTION
-from .extensions import ParameterTypeGuard
+from .extensions import ParameterTypeGuard, overload
 from .find_unused import UnusedObjectFinder, used
 from .reexport import ErrorContext, ImplicitReexportTracker
 from .safe import safe_getattr, is_hashable, safe_in, all_of_type
@@ -290,7 +291,9 @@ class _AttrContext(attributes.AttrContext):
 
 
 # FunctionInfo for a vanilla function (e.g. a lambda)
-_DEFAULT_FUNCTION_INFO = FunctionInfo(AsyncFunctionKind.normal, False, False, False, [])
+_DEFAULT_FUNCTION_INFO = FunctionInfo(
+    AsyncFunctionKind.normal, False, False, False, False, []
+)
 
 
 class ClassAttributeChecker:
@@ -1307,10 +1310,11 @@ class NameCheckVisitor(
             )
             expected_return_value = None
 
-        if evaled_function is not None:
-            self._set_name_in_scope(node.name, node, evaled_function)
-        else:
-            self._set_name_in_scope(node.name, node, KnownValue(potential_function))
+        if not info.is_overload:
+            if evaled_function is not None:
+                self._set_name_in_scope(node.name, node, evaled_function)
+            else:
+                self._set_name_in_scope(node.name, node, KnownValue(potential_function))
 
         with self.asynq_checker.set_func_name(
             node.name, async_kind=info.async_kind, is_classmethod=info.is_classmethod
@@ -1333,6 +1337,7 @@ class NameCheckVisitor(
 
         if (
             not has_return
+            and not info.is_overload
             and expected_return_value is not None
             and expected_return_value != KnownNone
             and not any(
@@ -1474,6 +1479,7 @@ class NameCheckVisitor(
         is_classmethod = False
         is_decorated_coroutine = False
         is_staticmethod = False
+        is_overload = False
         decorators = []
         for decorator in decorator_list:
             # We have to descend into the Call node because the result of
@@ -1501,12 +1507,17 @@ class NameCheckVisitor(
                     is_staticmethod = True
                 elif decorator_value == KnownValue(asyncio.coroutine):
                     is_decorated_coroutine = True
+                elif decorator_value == KnownValue(
+                    typing.overload
+                ) or decorator_value == KnownValue(overload):
+                    is_overload = True
                 decorators.append((decorator_value, decorator_value))
         return FunctionInfo(
             async_kind=async_kind,
             is_decorated_coroutine=is_decorated_coroutine,
             is_classmethod=is_classmethod,
             is_staticmethod=is_staticmethod,
+            is_overload=is_overload,
             decorators=decorators,
         )
 


Test patch:
diff --git a/pyanalyze/test_extensions.py b/pyanalyze/test_extensions.py
index d084e9d..25b63e6 100644
--- a/pyanalyze/test_extensions.py
+++ b/pyanalyze/test_extensions.py
@@ -1,8 +1,10 @@
 from qcore.asserts import AssertRaises
 import typing_inspect
 from typing import List, Optional, Union, TypeVar
+from types import FunctionType
 
-from .extensions import AsynqCallable
+from .extensions import AsynqCallable, get_overloads, overload
+from .safe import all_of_type
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -26,3 +28,48 @@ def test_asynq_callable() -> None:
         # get TypeVars out of an AsynqCallable instances. Solving this is hard
         # because Callable is special-cased at various places in typing.py.
         assert List[AsynqCallable[[str], int]] == List[AsynqCallable[[T], int]][str]
+
+
+@overload
+def f() -> int:
+    raise NotImplementedError
+
+
+@overload
+def f(a: int) -> str:
+    raise NotImplementedError
+
+
+def f(*args: object) -> object:
+    raise NotImplementedError
+
+
+class WithOverloadedMethods:
+    @overload
+    def f(self) -> int:
+        raise NotImplementedError
+
+    @overload
+    def f(self, a: int) -> str:
+        raise NotImplementedError
+
+    def f(self, *args: object) -> object:
+        raise NotImplementedError
+
+
+def test_overload() -> None:
+    overloads = get_overloads("pyanalyze.test_extensions.f")
+    assert len(overloads) == 2
+    assert all_of_type(overloads, FunctionType)
+    assert f not in overloads
+    assert overloads[0].__code__.co_argcount == 0
+    assert overloads[1].__code__.co_argcount == 1
+
+    method_overloads = get_overloads(
+        "pyanalyze.test_extensions.WithOverloadedMethods.f"
+    )
+    assert len(method_overloads) == 2
+    assert all_of_type(method_overloads, FunctionType)
+    assert WithOverloadedMethods.f not in overloads
+    assert method_overloads[0].__code__.co_argcount == 1
+    assert method_overloads[1].__code__.co_argcount == 2
diff --git a/pyanalyze/test_signature.py b/pyanalyze/test_signature.py
index 4d667be..662275f 100644
--- a/pyanalyze/test_signature.py
+++ b/pyanalyze/test_signature.py
@@ -1,4 +1,5 @@
 # static analysis: ignore
+from typed_ast.ast3 import Not
 from pyanalyze.implementation import assert_is_value
 from collections.abc import Sequence
 
@@ -992,7 +993,7 @@ class TestAnnotated(TestNameCheckVisitorBase):
 
 class TestOverload(TestNameCheckVisitorBase):
     @assert_passes()
-    def test(self):
+    def test_overloaded_impl(self):
         from pyanalyze.tests import overloaded
 
         def capybara():
@@ -1000,3 +1001,30 @@ class TestOverload(TestNameCheckVisitorBase):
             assert_is_value(overloaded("x"), TypedValue(str))
             overloaded(1)  # E: incompatible_call
             overloaded("x", "y")  # E: incompatible_call
+
+    @assert_passes()
+    def test_runtime(self):
+        from pyanalyze.extensions import overload
+        from typing import Union
+
+        @overload
+        def overloaded() -> int:
+            raise NotImplementedError
+
+        @overload
+        def overloaded(x: str) -> str:
+            raise NotImplementedError
+
+        def overloaded(*args: str) -> Union[int, str]:
+            if not args:
+                return 0
+            elif len(args) == 1:
+                return args[0]
+            else:
+                raise TypeError("too many arguments")
+
+        def capybara():
+            assert_is_value(overloaded(), TypedValue(int))
+            assert_is_value(overloaded("x"), TypedValue(str))
+            overloaded(1)  # E: incompatible_call
+            overloaded("a", "b")  # E: incompatible_call