Repository: quora/pyanalyze
Pull request #559 (2022-11-06 03:06:51)

More fixes to make coroutines return Coroutine not Awaitable


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 9249ba5..26d9abb 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -3,7 +3,7 @@
 ## Unreleased
 
 - Infer `async def` functions as returning `Coroutine`, not
-  `Awaitable` (#557)
+  `Awaitable` (#557, #559)
 - Drop support for Python 3.6 (#554)
 - Require `typeshed_client>=2.1.0`. Older versions will throw
   false-positive errors around context managers when
diff --git a/pyanalyze/annotations.py b/pyanalyze/annotations.py
index ecae666..ed6a1fb 100644
--- a/pyanalyze/annotations.py
+++ b/pyanalyze/annotations.py
@@ -76,7 +76,6 @@ from .signature import (
 )
 from .value import (
     _HashableValue,
-    TypeVarLike,
     annotate_value,
     AnnotatedValue,
     AnySource,
@@ -100,6 +99,7 @@ from .value import (
     TypedDictValue,
     TypedValue,
     TypeGuardExtension,
+    TypeVarLike,
     TypeVarValue,
     unite_values,
     UnpackedValue,
diff --git a/pyanalyze/arg_spec.py b/pyanalyze/arg_spec.py
index d01a5de..a2c0c38 100644
--- a/pyanalyze/arg_spec.py
+++ b/pyanalyze/arg_spec.py
@@ -11,7 +11,6 @@ import enum
 import inspect
 import sys
 import textwrap
-from collections.abc import Awaitable
 from dataclasses import dataclass, replace
 from types import FunctionType, MethodType, ModuleType
 from typing import (
@@ -75,6 +74,7 @@ from .value import (
     GenericValue,
     KnownValue,
     KVPair,
+    make_coro_type,
     NewTypeValue,
     SubclassValue,
     TypedDictValue,
@@ -344,7 +344,7 @@ class ArgSpecCache:
                 )
                 has_return_annotation = True
             if is_async:
-                returns = GenericValue(Awaitable, [returns])
+                returns = make_coro_type(returns)
 
         parameters = []
         for i, parameter in enumerate(sig.parameters.values()):
diff --git a/pyanalyze/functions.py b/pyanalyze/functions.py
index 95c1beb..a69dc0b 100644
--- a/pyanalyze/functions.py
+++ b/pyanalyze/functions.py
@@ -33,6 +33,7 @@ from .value import (
     GenericValue,
     get_tv_map,
     KnownValue,
+    make_coro_type,
     SubclassValue,
     TypedValue,
     TypeVarValue,
@@ -409,7 +410,7 @@ class IsGeneratorVisitor(ast.NodeVisitor):
     """Determine whether an async function is a generator.
 
     This is important because the return type of async generators
-    should not be wrapped in Awaitable.
+    should not be wrapped in Coroutine.
 
     We avoid recursing into nested functions, which is why we can't
     just use ast.walk.
@@ -449,10 +450,7 @@ def compute_value_of_function(
             if visitor.is_generator:
                 break
         if not visitor.is_generator:
-            result = GenericValue(
-                collections.abc.Coroutine,
-                [AnyValue(AnySource.inference), AnyValue(AnySource.inference), result],
-            )
+            result = make_coro_type(result)
     sig = Signature.make(
         [param_info.param for param_info in info.params],
         result,
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index a71ea6d..09b8875 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -173,6 +173,7 @@ from .value import (
     KnownValue,
     kv_pairs_from_mapping,
     KVPair,
+    make_coro_type,
     MultiValuedValue,
     NO_RETURN_VALUE,
     NoReturnConstraintExtension,
@@ -1801,7 +1802,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             return_value = AsyncTaskIncompleteValue(task_cls, return_value)
 
         if isinstance(info.node, ast.AsyncFunctionDef) or info.is_decorated_coroutine:
-            return_value = GenericValue(collections.abc.Awaitable, [return_value])
+            return_value = make_coro_type(return_value)
 
         if isinstance(val, KnownValue) and isinstance(val.val, property):
             fget = val.val.fget
diff --git a/pyanalyze/signature.py b/pyanalyze/signature.py
index 2a7d91c..83c3907 100644
--- a/pyanalyze/signature.py
+++ b/pyanalyze/signature.py
@@ -58,7 +58,6 @@ from .type_evaluation import (
 )
 from .typevar import resolve_bounds_map
 from .value import (
-    TypeVarLike,
     annotate_value,
     AnnotatedValue,
     AnySource,
@@ -95,6 +94,7 @@ from .value import (
     TypedDictValue,
     TypedValue,
     TypeGuardExtension,
+    TypeVarLike,
     TypeVarMap,
     TypeVarValue,
     unannotate,
diff --git a/pyanalyze/typeshed.py b/pyanalyze/typeshed.py
index 04e623e..d7f8ebf 100644
--- a/pyanalyze/typeshed.py
+++ b/pyanalyze/typeshed.py
@@ -11,13 +11,7 @@ import inspect
 import sys
 
 from abc import abstractmethod
-from collections.abc import (
-    Awaitable,
-    Collection,
-    MutableMapping,
-    Set as AbstractSet,
-    Sized,
-)
+from collections.abc import Collection, MutableMapping, Set as AbstractSet, Sized
 from dataclasses import dataclass, field, replace
 from enum import Enum, EnumMeta
 from types import GeneratorType, ModuleType
@@ -72,6 +66,7 @@ from .value import (
     extract_typevars,
     GenericValue,
     KnownValue,
+    make_coro_type,
     SubclassValue,
     TypedDictValue,
     TypedValue,
@@ -876,7 +871,7 @@ class TypeshedFinder:
         return Signature.make(
             cleaned_arguments,
             callable=obj,
-            return_annotation=GenericValue(Awaitable, [return_value])
+            return_annotation=make_coro_type(return_value)
             if isinstance(node, ast.AsyncFunctionDef)
             else return_value,
             allow_call=allow_call,
diff --git a/pyanalyze/value.py b/pyanalyze/value.py
index 0203e58..79f884c 100644
--- a/pyanalyze/value.py
+++ b/pyanalyze/value.py
@@ -2650,3 +2650,10 @@ def can_assign_and_used_any(
 
 def is_overlapping(left: Value, right: Value, ctx: CanAssignContext) -> bool:
     return left.is_assignable(right, ctx) or right.is_assignable(left, ctx)
+
+
+def make_coro_type(return_type: Value) -> GenericValue:
+    return GenericValue(
+        collections.abc.Coroutine,
+        [AnyValue(AnySource.inference), AnyValue(AnySource.inference), return_type],
+    )


Test patch:
diff --git a/pyanalyze/test.toml b/pyanalyze/test.toml
index a906a1f..e0537a0 100644
--- a/pyanalyze/test.toml
+++ b/pyanalyze/test.toml
@@ -6,4 +6,7 @@ constructor_hooks = ["pyanalyze.test_config.get_constructor"]
 known_signatures = ["pyanalyze.test_config.get_known_signatures"]
 unwrap_class = ["pyanalyze.test_config.unwrap_class"]
 stub_path = ["./stubs"]
-functions_safe_to_call = ["pyanalyze.tests.make_simple_sequence"]
+functions_safe_to_call = [
+    "pyanalyze.tests.make_simple_sequence",
+    "pyanalyze.value.make_coro_type",
+]
diff --git a/pyanalyze/test_async_await.py b/pyanalyze/test_async_await.py
index 72fb193..e6c6e0b 100644
--- a/pyanalyze/test_async_await.py
+++ b/pyanalyze/test_async_await.py
@@ -27,7 +27,7 @@ class TestAsyncDef(TestNameCheckVisitorBase):
 class TestAsyncAwait(TestNameCheckVisitorBase):
     @assert_passes()
     def test_type_inference(self):
-        from collections.abc import Awaitable
+        from pyanalyze.value import make_coro_type
 
         async def capybara(x):
             assert_is_value(x, AnyValue(AnySource.unannotated))
@@ -35,7 +35,7 @@ class TestAsyncAwait(TestNameCheckVisitorBase):
 
         async def kerodon(x):
             task = capybara(x)
-            assert_is_value(task, GenericValue(Awaitable, [KnownValue("hydrochoerus")]))
+            assert_is_value(task, make_coro_type(KnownValue("hydrochoerus")))
             val = await task
             assert_is_value(val, KnownValue("hydrochoerus"))
 
@@ -214,7 +214,8 @@ class TestArgSpec(TestNameCheckVisitorBase):
     @assert_passes()
     def test_asyncio_coroutine(self):
         import asyncio
-        from collections.abc import Awaitable
+
+        from pyanalyze.value import make_coro_type
 
         @asyncio.coroutine
         def f():
@@ -223,49 +224,45 @@ class TestArgSpec(TestNameCheckVisitorBase):
 
         @asyncio.coroutine
         def g():
-            assert_is_value(f(), GenericValue(Awaitable, [KnownValue(42)]))
+            assert_is_value(f(), make_coro_type(KnownValue(42)))
 
     @assert_passes()
     def test_coroutine_from_typeshed(self):
         import asyncio
-        import collections.abc
+
+        from pyanalyze.value import make_coro_type
 
         async def capybara():
-            assert_is_value(
-                asyncio.sleep(3),
-                GenericValue(collections.abc.Awaitable, [KnownValue(None)]),
-            )
+            assert_is_value(asyncio.sleep(3), make_coro_type(KnownValue(None)))
             return 42
 
     @assert_passes()
     def test_async_def_from_typeshed(self):
         from asyncio.streams import open_connection, StreamReader, StreamWriter
-        from collections.abc import Awaitable
+
+        from pyanalyze.value import make_coro_type
 
         async def capybara():
             # annotated as async def in typeshed
             assert_is_value(
                 open_connection(),
-                GenericValue(
-                    Awaitable,
-                    [
-                        make_simple_sequence(
-                            tuple, [TypedValue(StreamReader), TypedValue(StreamWriter)]
-                        )
-                    ],
+                make_coro_type(
+                    make_simple_sequence(
+                        tuple, [TypedValue(StreamReader), TypedValue(StreamWriter)]
+                    )
                 ),
             )
             return 42
 
     @assert_passes()
     def test_async_def(self):
-        from collections.abc import Awaitable
+        from pyanalyze.value import make_coro_type
 
         async def f():
             return 42
 
         async def g():
-            assert_is_value(f(), GenericValue(Awaitable, [KnownValue(42)]))
+            assert_is_value(f(), make_coro_type(KnownValue(42)))
 
 
 class TestNoReturn(TestNameCheckVisitorBase):
diff --git a/pyanalyze/test_attributes.py b/pyanalyze/test_attributes.py
index 05f1d4d..a1e5639 100644
--- a/pyanalyze/test_attributes.py
+++ b/pyanalyze/test_attributes.py
@@ -7,12 +7,12 @@ from .value import (
     AnnotatedValue,
     AnySource,
     AnyValue,
-    UnboundMethodValue,
     assert_is_value,
     GenericValue,
     KnownValue,
     MultiValuedValue,
     TypedValue,
+    UnboundMethodValue,
 )
 
 _global_dict: Dict[Union[int, str], float] = {}
@@ -173,11 +173,12 @@ class TestAttributes(TestNameCheckVisitorBase):
 
     @assert_passes()
     def test_annotated_known(self):
+        from qcore.testing import Anything
         from typing_extensions import Annotated, Literal
+
         from pyanalyze.extensions import LiteralOnly
         from pyanalyze.stacked_scopes import Composite, VarnameWithOrigin
         from pyanalyze.value import CustomCheckExtension
-        from qcore.testing import Anything
 
         origin = VarnameWithOrigin("encoding", Anything)  # E: incompatible_argument
 
diff --git a/pyanalyze/test_functions.py b/pyanalyze/test_functions.py
index 1987657..9da0564 100644
--- a/pyanalyze/test_functions.py
+++ b/pyanalyze/test_functions.py
@@ -51,23 +51,13 @@ class TestNestedFunction(TestNameCheckVisitorBase):
 
     @assert_passes()
     def test_async_def(self):
-        import collections.abc
+        from pyanalyze.value import make_coro_type
 
         def capybara():
             async def nested() -> int:
                 return 1
 
-            assert_is_value(
-                nested(),
-                GenericValue(
-                    collections.abc.Coroutine,
-                    [
-                        AnyValue(AnySource.inference),
-                        AnyValue(AnySource.inference),
-                        TypedValue(int),
-                    ],
-                ),
-            )
+            assert_is_value(nested(), make_coro_type(TypedValue(int)))
 
     @assert_passes()
     def test_bad_decorator(self):
@@ -198,7 +188,7 @@ class TestDecorators(TestNameCheckVisitorBase):
 class TestAsyncGenerator(TestNameCheckVisitorBase):
     @assert_passes()
     def test_not_a_generator(self):
-        from collections.abc import Awaitable
+        from pyanalyze.value import make_coro_type
 
         async def capybara() -> None:
             async def agen():
@@ -211,7 +201,7 @@ class TestAsyncGenerator(TestNameCheckVisitorBase):
 
         def caller() -> None:
             x = capybara()
-            assert_is_value(x, GenericValue(Awaitable, [KnownValue(None)]))
+            assert_is_value(x, make_coro_type(KnownValue(None)))
 
     @assert_passes()
     def test_is_a_generator(self):
diff --git a/pyanalyze/test_signature.py b/pyanalyze/test_signature.py
index 8d483ea..96c3ac3 100644
--- a/pyanalyze/test_signature.py
+++ b/pyanalyze/test_signature.py
@@ -887,7 +887,8 @@ class TestAllowCall(TestNameCheckVisitorBase):
 
     @assert_passes()
     def test_annotated_known(self):
-        from typing_extensions import Literal, Annotated
+        from typing_extensions import Annotated, Literal
+
         from pyanalyze.extensions import LiteralOnly
 
         def capybara():
diff --git a/pyanalyze/test_typevar.py b/pyanalyze/test_typevar.py
index 8be2978..1bdb3de 100644
--- a/pyanalyze/test_typevar.py
+++ b/pyanalyze/test_typevar.py
@@ -278,6 +278,7 @@ class TestSolve(TestNameCheckVisitorBase):
     @assert_passes()
     def test_redundant_constraints(self):
         from typing import TypeVar
+
         from typing_extensions import SupportsIndex
 
         T = TypeVar("T", int, float, SupportsIndex)
@@ -293,6 +294,7 @@ class TestSolve(TestNameCheckVisitorBase):
     @assert_passes()
     def test_lots_of_constraints(self):
         from typing import TypeVar, Union
+
         from typing_extensions import SupportsIndex
 
         T = TypeVar(
diff --git a/pyanalyze/test_yield_checker.py b/pyanalyze/test_yield_checker.py
index 0755bc0..d557ee2 100644
--- a/pyanalyze/test_yield_checker.py
+++ b/pyanalyze/test_yield_checker.py
@@ -501,6 +501,7 @@ class TestBatchingYields(TestNameCheckVisitorBase):
     @assert_passes()
     def test_combine_in_try(self):
         from asynq import asynq
+
         from pyanalyze.tests import async_fn
 
         @asynq()