Repository: quora/pyanalyze
Pull request #325 (2021-12-08 16:32:52)

support overloads from typeshed

JelleZijlstra: Notes:
* I found while implementing this that the stub parser doesn't support `|` in types in stubs, because of the ast3/ast difference. This should be fixed in its own PR because it could likely find a lot of errors by itself.
* I'm seeing internal errors on `int **` because the overloads for that use Literals of negative numbers, and the type-from-AST visitor doesn't support UnaryOp. This should also be fixed separately.
* Lots of errors due to incompatible types with dict.update. We should add an impl function for it that can widen types.
JelleZijlstra: (1) is especially bad, because it silently made a lot of types Any. I should add a test that iterates over all the typeshed stubs and asserts that we can parse them. This could eventually even go into typeshed.
JelleZijlstra: There is also an issue with the bases for `urllib.parse.ParseResult` that shows up when you call `urllib.parse.urlparse`.

Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index ad32005..2a27450 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Recognize overloaded functions in stubs (#325)
 - Fix handling of classes in stubs that have an incorrect `__qualname__`
   at runtime (#336)
 - Fix type compatibility with generic functions (#335)
diff --git a/pyanalyze/annotations.py b/pyanalyze/annotations.py
index b742a2e..3c85034 100644
--- a/pyanalyze/annotations.py
+++ b/pyanalyze/annotations.py
@@ -31,6 +31,7 @@ import ast
 from typed_ast import ast3
 import builtins
 from collections.abc import Callable, Iterable
+from typed_ast import ast3
 from typing import (
     Any,
     Container,
diff --git a/pyanalyze/arg_spec.py b/pyanalyze/arg_spec.py
index 8f0384d..d9c02c5 100644
--- a/pyanalyze/arg_spec.py
+++ b/pyanalyze/arg_spec.py
@@ -386,7 +386,8 @@ class ArgSpecCache:
                     original_fn, impl, is_asynq, in_overload_resolution
                 )
 
-        argspec = self.ts_finder.get_argspec(obj)
+        allow_call = safe_in(obj, self.config.FUNCTIONS_SAFE_TO_CALL)
+        argspec = self.ts_finder.get_argspec(obj, allow_call=allow_call)
         if argspec is not None:
             return argspec
 
diff --git a/pyanalyze/implementation.py b/pyanalyze/implementation.py
index c8b405b..cdbde82 100644
--- a/pyanalyze/implementation.py
+++ b/pyanalyze/implementation.py
@@ -1420,6 +1420,42 @@ def get_default_argspecs() -> Dict[object, Signature]:
             callable=len,
             impl=_len_impl,
         ),
+        # The overloaded annotation in typeshed causes a couple of problems:
+        # - sorted(Sequence[A] | Sequence[B]) turns into List[A | B] instead of List[A] | List[B]
+        #   pyright agrees with pyanalyze here but mypy somehow infers the second type.
+        # - The bounded TypeVar on the key argument makes pyanalyze infer the
+        #   return type as list[Sized] if we use key=len. Fixing this may require changing the
+        #   TypeVar resolution algorithm.
+        # To avoid these problems, we use a more permissive hardcoded signature for now.
+        Signature.make(
+            [
+                SigParameter(
+                    "iterable",
+                    SigParameter.POSITIONAL_ONLY,
+                    annotation=TypedValue(collections.abc.Iterable),
+                ),
+                SigParameter(
+                    "key",
+                    SigParameter.KEYWORD_ONLY,
+                    annotation=CallableValue(
+                        Signature.make(
+                            [SigParameter("arg", SigParameter.POSITIONAL_ONLY)],
+                            return_annotation=TypedValue("_typeshed.SupportsLessThan"),
+                        )
+                    ),
+                    default=KnownValue(None),
+                ),
+                SigParameter(
+                    "reverse",
+                    SigParameter.KEYWORD_ONLY,
+                    annotation=TypedValue(bool),
+                    default=KnownValue(False),
+                ),
+            ],
+            return_annotation=TypedValue(list),
+            allow_call=True,
+            callable=sorted,
+        ),
         # TypeGuards, which aren't in typeshed yet
         Signature.make(
             [
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 70c8ee0..601a1cc 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -2264,10 +2264,12 @@ class NameCheckVisitor(
                     keywords = {kw.arg for kw in parent.keywords if kw.arg is not None}
                     if names <= keywords:
                         return
+            stmt = f_str_ast.body[0]
+            assert isinstance(stmt, ast.Expr), f"unexpected ast {ast.dump(f_str_ast)}"
             self._show_error_if_checking(
                 node,
                 error_code=ErrorCode.missing_f,
-                replacement=self.replace_node(node, f_str_ast.body[0].value),
+                replacement=self.replace_node(node, stmt.value),
             )
 
     def _name_exists(self, name: str) -> bool:
diff --git a/pyanalyze/node_visitor.py b/pyanalyze/node_visitor.py
index ef59ae9..cc6b861 100644
--- a/pyanalyze/node_visitor.py
+++ b/pyanalyze/node_visitor.py
@@ -473,10 +473,11 @@ class BaseNodeVisitor(ast.NodeVisitor):
 
     @classmethod
     def _apply_changes_to_lines(
-        cls, changes: List[Replacement], lines: Sequence[str]
+        cls, changes: List[Replacement], input_lines: Sequence[str]
     ) -> Sequence[str]:
         # only apply the first change because that change might affect other fixes
         # that test_scope came up for that file. So we break after finding first applicable fix.
+        lines = list(input_lines)
         if changes:
             change = changes[0]
             additions = change.lines_to_add
@@ -484,7 +485,7 @@ class BaseNodeVisitor(ast.NodeVisitor):
                 lines_to_remove = change.linenos_to_delete
                 max_line = max(lines_to_remove)
                 # add the additions after the max_line
-                lines = lines[:max_line] + additions + lines[max_line:]
+                lines = [*lines[:max_line], *additions, *lines[max_line:]]
                 lines_to_remove = sorted(lines_to_remove, reverse=True)
                 for lineno in lines_to_remove:
                     del lines[lineno - 1]
diff --git a/pyanalyze/signature.py b/pyanalyze/signature.py
index a116928..b8a5a39 100644
--- a/pyanalyze/signature.py
+++ b/pyanalyze/signature.py
@@ -1106,17 +1106,17 @@ class Signature:
             elif param.kind is SigParameter.POSITIONAL_ONLY:
                 if param.name not in consumed_positional:
                     return CanAssignError(
-                        f"takes extra positional-only paramter {param.name!r}"
+                        f"takes extra positional-only parameter {param.name!r}"
                     )
             elif param.kind is SigParameter.POSITIONAL_OR_KEYWORD:
                 if (
                     param.name not in consumed_positional
                     and param.name not in consumed_keyword
                 ):
-                    return CanAssignError(f"takes extra paramter {param.name!r}")
+                    return CanAssignError(f"takes extra parameter {param.name!r}")
             elif param.kind is SigParameter.KEYWORD_ONLY:
                 if param.name not in consumed_keyword:
-                    return CanAssignError(f"takes extra paramter {param.name!r}")
+                    return CanAssignError(f"takes extra parameter {param.name!r}")
 
         return unify_typevar_maps(tv_maps)
 
diff --git a/pyanalyze/typeshed.py b/pyanalyze/typeshed.py
index 9b19c73..b76a8ca 100644
--- a/pyanalyze/typeshed.py
+++ b/pyanalyze/typeshed.py
@@ -8,7 +8,7 @@ from .annotations import Context, type_from_value, value_from_ast
 from .error_code import ErrorCode
 from .safe import is_typing_name
 from .stacked_scopes import uniq_chain
-from .signature import SigParameter, Signature
+from .signature import ConcreteSignature, OverloadedSignature, SigParameter, Signature
 from .value import (
     AnySource,
     AnyValue,
@@ -52,6 +52,7 @@ from typing import (
     Callable,
     List,
     TypeVar,
+    overload,
 )
 from typing_extensions import Protocol, TypedDict
 import typeshed_client
@@ -110,7 +111,9 @@ class TypeshedFinder:
             return
         print("%s: %r" % (message, obj))
 
-    def get_argspec(self, obj: object) -> Optional[Signature]:
+    def get_argspec(
+        self, obj: object, *, allow_call: bool = False
+    ) -> Optional[ConcreteSignature]:
         if inspect.ismethoddescriptor(obj) and hasattr(obj, "__objclass__"):
             objclass = obj.__objclass__
             fq_name = self._get_fq_name(objclass)
@@ -118,7 +121,7 @@ class TypeshedFinder:
                 return None
             info = self._get_info_for_name(fq_name)
             sig = self._get_method_signature_from_info(
-                info, obj, fq_name, objclass.__module__, objclass
+                info, obj, fq_name, objclass.__module__, objclass, allow_call=allow_call
             )
             return sig
 
@@ -140,20 +143,26 @@ class TypeshedFinder:
                 if maybe_info is not None:
                     info, mod = maybe_info
                     fq_name = f"{parent_fqn}.{own_name}"
-                    sig = self._get_signature_from_info(info, obj, fq_name, mod)
+                    sig = self._get_signature_from_info(
+                        info, obj, fq_name, mod, allow_call=allow_call
+                    )
                     return sig
 
         fq_name = self._get_fq_name(obj)
         if fq_name is None:
             return None
-        return self.get_argspec_for_fully_qualified_name(fq_name, obj)
+        return self.get_argspec_for_fully_qualified_name(
+            fq_name, obj, allow_call=allow_call
+        )
 
     def get_argspec_for_fully_qualified_name(
-        self, fq_name: str, obj: object
-    ) -> Optional[Signature]:
+        self, fq_name: str, obj: object, *, allow_call: bool = False
+    ) -> Optional[ConcreteSignature]:
         info = self._get_info_for_name(fq_name)
         mod, _ = fq_name.rsplit(".", maxsplit=1)
-        sig = self._get_signature_from_info(info, obj, fq_name, mod)
+        sig = self._get_signature_from_info(
+            info, obj, fq_name, mod, allow_call=allow_call
+        )
         return sig
 
     def get_bases(self, typ: type) -> Optional[List[Value]]:
@@ -498,19 +507,26 @@ class TypeshedFinder:
         fq_name: str,
         mod: str,
         objclass: type,
-    ) -> Optional[Signature]:
+        *,
+        allow_call: bool = False,
+    ) -> Optional[ConcreteSignature]:
         if info is None:
             return None
         elif isinstance(info, typeshed_client.ImportedInfo):
             return self._get_method_signature_from_info(
-                info.info, obj, fq_name, ".".join(info.source_module), objclass
+                info.info,
+                obj,
+                fq_name,
+                ".".join(info.source_module),
+                objclass,
+                allow_call=allow_call,
             )
         elif isinstance(info, typeshed_client.NameInfo):
             # Note that this doesn't handle names inherited from base classes
             if info.child_nodes and obj.__name__ in info.child_nodes:
                 child_info = info.child_nodes[obj.__name__]
                 return self._get_signature_from_info(
-                    child_info, obj, fq_name, mod, objclass
+                    child_info, obj, fq_name, mod, objclass, allow_call=allow_call
                 )
             else:
                 return None
@@ -550,16 +566,41 @@ class TypeshedFinder:
         fq_name: str,
         mod: str,
         objclass: Optional[type] = None,
-    ) -> Optional[Signature]:
+        *,
+        allow_call: bool = False,
+    ) -> Optional[ConcreteSignature]:
         if isinstance(info, typeshed_client.NameInfo):
             if isinstance(info.ast, (ast3.FunctionDef, ast3.AsyncFunctionDef)):
-                return self._get_signature_from_func_def(info.ast, obj, mod, objclass)
+                return self._get_signature_from_func_def(
+                    info.ast, obj, mod, objclass, allow_call=allow_call
+                )
+            elif isinstance(info.ast, typeshed_client.OverloadedName):
+                sigs = []
+                for defn in info.ast.definitions:
+                    if not isinstance(defn, (ast3.FunctionDef, ast3.AsyncFunctionDef)):
+                        self.log(
+                            "Ignoring unrecognized AST in overload", (fq_name, info)
+                        )
+                        return None
+                    sig = self._get_signature_from_func_def(
+                        defn, obj, mod, objclass, allow_call=allow_call
+                    )
+                    if sig is None:
+                        self.log("Could not get sig for overload member", (defn,))
+                        return None
+                    sigs.append(sig)
+                return OverloadedSignature(sigs)
             else:
                 self.log("Ignoring unrecognized AST", (fq_name, info))
                 return None
         elif isinstance(info, typeshed_client.ImportedInfo):
             return self._get_signature_from_info(
-                info.info, obj, fq_name, ".".join(info.source_module), objclass
+                info.info,
+                obj,
+                fq_name,
+                ".".join(info.source_module),
+                objclass,
+                allow_call=allow_call,
             )
         elif info is None:
             return None
@@ -579,11 +620,14 @@ class TypeshedFinder:
         objclass: Optional[type] = None,
         *,
         autobind: bool = False,
+        allow_call: bool = False,
     ) -> Optional[Signature]:
         is_classmethod = is_staticmethod = False
         for decorator_ast in node.decorator_list:
             decorator = self._parse_expr(decorator_ast, mod)
-            if decorator == KnownValue(abstractmethod):
+            if decorator == KnownValue(abstractmethod) or decorator == KnownValue(
+                overload
+            ):
                 continue
             elif decorator == KnownValue(classmethod):
                 is_classmethod = True
@@ -650,6 +694,7 @@ class TypeshedFinder:
             return_annotation=GenericValue(Awaitable, [return_value])
             if isinstance(node, ast3.AsyncFunctionDef)
             else return_value,
+            allow_call=allow_call,
         )
 
     def _parse_param_list(
@@ -699,10 +744,10 @@ class TypeshedFinder:
         if default is None:
             return SigParameter(name, kind, annotation=typ)
         else:
-            default = self._parse_expr(default, module)
-            if default == KnownValue(...):
-                default = AnyValue(AnySource.unannotated)
-            return SigParameter(name, kind, annotation=typ, default=default)
+            default_value = self._parse_expr(default, module)
+            if default_value == KnownValue(...):
+                default_value = AnyValue(AnySource.unannotated)
+            return SigParameter(name, kind, annotation=typ, default=default_value)
 
     def _parse_expr(self, node: ast3.AST, module: str) -> Value:
         ctx = _AnnotationContext(finder=self, module=module)
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index 9347d5c..afed0a5 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -579,8 +579,10 @@ class TypedValue(Value):
         elif isinstance(other, TypedValue):
             return self_tobj.can_assign(self, other, ctx)
         elif isinstance(other, SubclassValue):
-            if isinstance(other.typ, TypedValue) and isinstance(
-                other.typ.typ, self.typ
+            if (
+                isinstance(other.typ, TypedValue)
+                and isinstance(self.typ, type)
+                and isinstance(other.typ.typ, self.typ)
             ):
                 return {}
             elif isinstance(other.typ, (TypeVarValue, AnyValue)):


Test patch:
diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
index d188178..eba77ab 100644
--- a/pyanalyze/test_name_check_visitor.py
+++ b/pyanalyze/test_name_check_visitor.py
@@ -1050,8 +1050,8 @@ class TestClassAttributeChecker(TestNameCheckVisitorBase):
     @assert_passes()
     def test_setattr(self):
         class Capybara(object):
-            def __init__(self):
-                for k, v in iter([("grass", "tasty")]):
+            def __init__(self, unannotated):
+                for k, v in unannotated:
                     assert_is_value(k, AnyValue(AnySource.generic_argument))
                     setattr(self, k, v)
                 assert_is_value(self.grass, AnyValue(AnySource.inference))
@@ -1059,8 +1059,8 @@ class TestClassAttributeChecker(TestNameCheckVisitorBase):
     @assert_passes()
     def test_setattr_on_base(self):
         class Capybara(object):
-            def __init__(self):
-                for k, v in iter([("grass", "tasty")]):
+            def __init__(self, unannotated):
+                for k, v in unannotated:
                     # Make sure we're not smart enough to infer the attribute
                     assert_is_value(k, AnyValue(AnySource.generic_argument))
                     setattr(self, k, v)
@@ -1782,10 +1782,15 @@ class TestOperators(TestNameCheckVisitorBase):
             assert_is_value(ha + x, AnyValue(AnySource.from_another))
             assert_is_value(HasBoth() + HasBoth(), TypedValue(HasBoth))
 
-    @assert_fails(ErrorCode.unsupported_operation)
+    @assert_passes()
     def test_unsupported_unary_op(self):
         def capybara():
-            ~"capybara"
+            ~"capybara"  # E: unsupported_operation
+
+    @assert_passes()
+    def test_int_float_product(self):
+        def capybara(f: float, i: int):
+            assert_is_value(i * f, TypedValue(float))
 
 
 class TestTaskNeedsYield(TestNameCheckVisitorBase):
diff --git a/pyanalyze/test_signature.py b/pyanalyze/test_signature.py
index 0562fd0..d5506a7 100644
--- a/pyanalyze/test_signature.py
+++ b/pyanalyze/test_signature.py
@@ -1106,3 +1106,27 @@ class TestOverload(TestNameCheckVisitorBase):
             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))
+
+    @assert_passes()
+    def test_typeshed_overload(self):
+        class SupportsWrite:
+            def write(self, s: str) -> None:
+                pass
+
+        class SupportsWriteAndFlush(SupportsWrite):
+            def flush(self) -> None:
+                pass
+
+        def capybara():
+            print()  # ok
+            print("x", file=SupportsWrite())
+            print("x", file=SupportsWrite(), flush=True)  # E: incompatible_argument
+            print("x", file=SupportsWriteAndFlush(), flush=True)
+            print("x", file=SupportsWriteAndFlush())
+            print("x", file="not a file")  # E: incompatible_call
+
+        def pacarana(f: float):
+            assert_is_value(f.__round__(), TypedValue(int))
+            assert_is_value(f.__round__(None), TypedValue(int))
+            f.__round__(ndigits=None)  # E: incompatible_call
+            assert_is_value(f.__round__(1), TypedValue(float))