Repository: quora/pyanalyze (python)

Issue:
Support overloaded functions
Pyanalyze currently doesn't support overloaded functions at all. It should. Some thoughts in no particular order:
- Supporting overloads for runtime code is hard because the overloads aren't stored in the module dict. The best solution I can think of is to add a `pyanalyze.extensions.overload` decorator that somehow does register the overloads at runtime. It could be an alias of `from typing import overload` under `if TYPE_CHECKING` for compatibility with other type checkers. We could add an `@overload_implementation` decorator that collects the other overloads and stores them in the function dict or something.
- But it's much easier to support them in stubs and `implementation.py`. It might require changes to typeshed-client though.
- Internally an overloaded function could be represented as a new sister class of `pyanalyze.signature.Signature`. It would hold all of the overloaded signatures in an attribute.
- Checking a call to an overload should try them one by one and succeed if any overload succeeds. But consider this example:
```
@overload
def f(x: int) -> Literal[1]: ...
@overload
def f(x: str) -> Literal[2]: ...
y: Union[int, str]
f(y)
```
I would want this to return `Literal[1, 2]`, but naively looping over the overloads would not work for that; it would require special handling for Union. Interestingly, though, TypeScript doesn't let you do this. You have to write a separate overload for the Union instead.

- Overloads could be useful internally for representing the generic mutators (`list.append` etc.), for which there are currently complicated impl functions. These overloads would require being able to dispatch on whether a value is annotated with WeakExtension.

Pull request description:
patch @typing.overload
So we can (usually) typecheck libraries that use `@overload`.

Closes #182.

Patch:
diff --git a/pyanalyze/arg_spec.py b/pyanalyze/arg_spec.py
index 2448a10..3d5a0c6 100644
--- a/pyanalyze/arg_spec.py
+++ b/pyanalyze/arg_spec.py
@@ -18,6 +18,7 @@ from .safe import (
     safe_issubclass,
     is_typing_name,
     safe_isinstance,
+    get_fully_qualified_name,
 )
 from .stacked_scopes import Composite, uniq_chain
 from .signature import (
@@ -516,7 +517,7 @@ class ArgSpecCache:
         in_overload_resolution: bool,
     ) -> MaybeSignature:
         if not in_overload_resolution:
-            fq_name = _get_fully_qualified_name(obj)
+            fq_name = get_fully_qualified_name(obj)
             if fq_name is not None:
                 overloads = get_overloads(fq_name)
                 if overloads:
@@ -848,12 +849,6 @@ def _is_qcore_decorator(obj: object) -> bool:
         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
-
-
 def _get_class_name(obj: object) -> Optional[str]:
     if hasattr(obj, "__qualname__"):
         pieces = obj.__qualname__.split(".")
diff --git a/pyanalyze/extensions.py b/pyanalyze/extensions.py
index 7b6643d..b684cff 100644
--- a/pyanalyze/extensions.py
+++ b/pyanalyze/extensions.py
@@ -10,6 +10,7 @@ be gracefully ignored by other type checkers.
 """
 from collections import defaultdict
 from dataclasses import dataclass, field
+import typing_extensions
 import pyanalyze
 from typing import (
     Any,
@@ -26,6 +27,9 @@ from typing import (
     overload as real_overload,
 )
 from typing_extensions import Literal, NoReturn
+import typing
+
+from .safe import get_fully_qualified_name
 
 if TYPE_CHECKING:
     from .value import Value, CanAssign, CanAssignContext, TypeVarMap, AnySource
@@ -409,11 +413,23 @@ else:
         return all the runtime overloads.
 
         """
-        key = f"{func.__module__}.{func.__qualname__}"
-        _overloads[key].append(func)
+        key = get_fully_qualified_name(func)
+        if key is not None:
+            _overloads[key].append(func)
         return real_overload(func)
 
 
+def patch_typing_overload() -> None:
+    """Monkey-patch ``typing.overload`` with our custom ``@overload`` decorator.
+
+    This allows files imported after this file to use the ``@overload`` decorator
+    and have it be recognized by pyanalyze.
+
+    """
+    typing.overload = overload
+    typing_extensions.overload = overload
+
+
 def evaluated(func: Callable[..., Any]) -> Callable[..., Any]:
     """Marks a type evaluation function."""
     key = f"{func.__module__}.{func.__qualname__}"
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 9b142e5..c513d00 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -68,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, patch_typing_overload
 from .find_unused import UnusedObjectFinder, used
 from .options import (
     ConfigOption,
@@ -4642,6 +4642,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             options.display()
             sys.exit(0)
         kwargs.setdefault("checker", Checker(cls.config, options))
+        patch_typing_overload()
         return kwargs
 
     def is_enabled(self, error_code: enum.Enum) -> bool:
diff --git a/pyanalyze/safe.py b/pyanalyze/safe.py
index 889350d..14f1d7f 100644
--- a/pyanalyze/safe.py
+++ b/pyanalyze/safe.py
@@ -5,7 +5,18 @@
 """
 import inspect
 import sys
-from typing import Any, Dict, Sequence, Tuple, Union, Container, NewType, Type, TypeVar
+from typing import (
+    Any,
+    Dict,
+    Sequence,
+    Tuple,
+    Union,
+    Container,
+    NewType,
+    Type,
+    TypeVar,
+    Optional,
+)
 import typing
 import typing_extensions
 
@@ -152,3 +163,9 @@ def _fill_typing_name_cache(name: str) -> Tuple[Tuple[Any, ...], Tuple[str, ...]
         result = tuple(objs), tuple(names)
         _typing_name_cache[name] = result
         return result
+
+
+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/typeshed.py b/pyanalyze/typeshed.py
index b8b77c1..b86262c 100644
--- a/pyanalyze/typeshed.py
+++ b/pyanalyze/typeshed.py
@@ -16,6 +16,7 @@ from .annotations import (
     value_from_ast,
 )
 from .error_code import ErrorCode
+from .extensions import overload, real_overload
 from .safe import is_typing_name
 from .stacked_scopes import uniq_chain
 from .signature import (
@@ -71,7 +72,6 @@ from typing import (
     Callable,
     List,
     TypeVar,
-    overload,
 )
 from typing_extensions import Protocol, TypedDict
 import typeshed_client
@@ -691,8 +691,10 @@ class TypeshedFinder:
         is_classmethod = is_staticmethod = is_evaluated = False
         for decorator_ast in node.decorator_list:
             decorator = self._parse_expr(decorator_ast, mod)
-            if decorator == KnownValue(abstractmethod) or decorator == KnownValue(
-                overload
+            if (
+                decorator == KnownValue(abstractmethod)
+                or decorator == KnownValue(overload)
+                or decorator == KnownValue(real_overload)
             ):
                 continue
             elif decorator == KnownValue(classmethod):


Test patch:
diff --git a/pyanalyze/test_signature.py b/pyanalyze/test_signature.py
index a998aaf..26eccf5 100644
--- a/pyanalyze/test_signature.py
+++ b/pyanalyze/test_signature.py
@@ -1150,3 +1150,30 @@ class TestOverload(TestNameCheckVisitorBase):
             assert_is_value(f.__round__(None), TypedValue(int))
             f.__round__(ndigits=None)  # E: incompatible_call
             assert_is_value(f.__round__(1), TypedValue(float))
+
+    @assert_passes()
+    def test_runtime(self):
+        from pyanalyze.extensions import (
+            overload as extension_overload,
+            patch_typing_overload,
+        )
+
+        patch_typing_overload()
+        from typing import overload
+
+        assert overload is extension_overload, (overload, extension_overload)
+
+        @overload
+        def f(x: int) -> str:
+            pass
+
+        @overload
+        def f(x: str) -> int:
+            pass
+
+        def f(x: object) -> object:
+            raise NotImplementedError
+
+        def capybara():
+            assert_is_value(f(1), TypedValue(str))
+            assert_is_value(f(""), TypedValue(int))