[ty] Support async/await

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/async.md

@@ -0,0 +1,15 @@
+# `async` / `await`
+
+## Basic
+
+```py
+async def retrieve() -> int:
+    return 42
+
+reveal_type(retrieve)  # revealed: def retrieve() -> CoroutineType[Any, Any, int]
+
+async def main():
+    result = await retrieve()
+
+    reveal_type(result)  # revealed: int
+```

crates/ty_python_semantic/resources/mdtest/call/function.md

@@ -15,8 +15,7 @@ reveal_type(get_int())  # revealed: int
 async def get_int_async() -> int:
     return 42
 
-# TODO: we don't yet support `types.CoroutineType`, should be generic `Coroutine[Any, Any, int]`
-reveal_type(get_int_async())  # revealed: @Todo(generic types.CoroutineType)
+reveal_type(get_int_async())  # revealed: CoroutineType[Any, Any, int]
 ```
 
 ## Generic

crates/ty_python_semantic/resources/mdtest/function/return_type.md

@@ -433,6 +433,8 @@ def f(cond: bool) -> str:
 
 <!-- snapshot-diagnostics -->
 
+## Synchronous
+
 A function with a `yield` or `yield from` expression anywhere in its body is a
 [generator function](https://docs.python.org/3/glossary.html#term-generator). A generator function
 implicitly returns an instance of `types.GeneratorType` even if it does not contain any `return`
@@ -461,6 +463,8 @@ def j() -> str:  # error: [invalid-return-type]
     yield 42
 ```
 
+## Asynchronous
+
 If it is an `async` function with a `yield` statement in its body, it is an
 [asynchronous generator function](https://docs.python.org/3/glossary.html#term-asynchronous-generator).
 An asynchronous generator function implicitly returns an instance of `types.AsyncGeneratorType` even

crates/ty_python_semantic/resources/mdtest/snapshots/return_type.md_-_Function_return_type_-_Generator_functions_(d9ed06b61b14fd4c).snap

@@ -28,7 +28,7 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/function/return_type.md
 14 |     yield 42
 15 | 
 16 | def i2() -> typing.Generator:
-17 |     yield from i()
+17 |     yield from reveal_type(i())  # error: [not-iterable]
 18 | 
 19 | def j() -> str:  # error: [invalid-return-type]
 20 |     yield 42
@@ -53,11 +53,54 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/function/return_type.md
 
 # Diagnostics
 
+```
+error[not-iterable]: Object of type `Iterable[Unknown] | CoroutineType[Any, Any, AsyncIterable[Unknown]]` may not be iterable
+  --> src/mdtest_snippet.py:17:16
+   |
+16 | def i2() -> typing.Generator:
+17 |     yield from reveal_type(i())  # error: [not-iterable]
+   |                ^^^^^^^^^^^^^^^^
+18 |
+19 | def j() -> str:  # error: [invalid-return-type]
+   |
+info: It may not have an `__iter__` method and it doesn't have a `__getitem__` method
+info: rule `not-iterable` is enabled by default
+
+```
+
+```
+warning[undefined-reveal]: `reveal_type` used without importing it
+  --> src/mdtest_snippet.py:17:16
+   |
+16 | def i2() -> typing.Generator:
+17 |     yield from reveal_type(i())  # error: [not-iterable]
+   |                ^^^^^^^^^^^
+18 |
+19 | def j() -> str:  # error: [invalid-return-type]
+   |
+info: This is allowed for debugging convenience but will fail at runtime
+info: rule `undefined-reveal` is enabled by default
+
+```
+
+```
+info[revealed-type]: Revealed type
+  --> src/mdtest_snippet.py:17:28
+   |
+16 | def i2() -> typing.Generator:
+17 |     yield from reveal_type(i())  # error: [not-iterable]
+   |                            ^^^ `Iterable[Unknown] | CoroutineType[Any, Any, AsyncIterable[Unknown]]`
+18 |
+19 | def j() -> str:  # error: [invalid-return-type]
+   |
+
+```
+
 ```
 error[invalid-return-type]: Return type does not match returned value
   --> src/mdtest_snippet.py:19:12
    |
-17 |     yield from i()
+17 |     yield from reveal_type(i())  # error: [not-iterable]
 18 |
 19 | def j() -> str:  # error: [invalid-return-type]
    |            ^^^ expected `str`, found `types.GeneratorType`

crates/ty_python_semantic/src/types/class.rs

@@ -2610,10 +2610,13 @@ pub enum KnownClass {
     UnionType,
     GeneratorType,
     AsyncGeneratorType,
+    CoroutineType,
     // Typeshed
     NoneType, // Part of `types` for Python >= 3.10
     // Typing
     Any,
+    Awaitable,
+    Generator,
     Deprecated,
     StdlibAlias,
     SpecialForm,
@@ -2736,6 +2739,9 @@ impl KnownClass {
             | Self::NotImplementedType
             | Self::Staticmethod
             | Self::Classmethod
+            | Self::Awaitable
+            | Self::Generator
+            | Self::CoroutineType
             | Self::Deprecated
             | Self::Field
             | Self::KwOnly
@@ -2771,7 +2777,8 @@ impl KnownClass {
             | Self::Super
             | Self::GenericAlias
             | Self::Deque
-            | Self::Bytes => true,
+            | Self::Bytes
+            | Self::CoroutineType => true,
 
             // It doesn't really make sense to ask the question for `@final` types,
             // since these are "more than solid bases". But we'll anyway infer a `@final`
@@ -2807,6 +2814,8 @@ impl KnownClass {
             // with length >2, or anything that is implemented in pure Python, is not a solid base.
             Self::ABCMeta
             | Self::Any
+            | Self::Awaitable
+            | Self::Generator
             | Self::Enum
             | Self::EnumType
             | Self::Auto
@@ -2855,6 +2864,8 @@ impl KnownClass {
             | KnownClass::ExceptionGroup
             | KnownClass::Staticmethod
             | KnownClass::Classmethod
+            | KnownClass::Awaitable
+            | KnownClass::Generator
             | KnownClass::Deprecated
             | KnownClass::Super
             | KnownClass::Enum
@@ -2872,6 +2883,7 @@ impl KnownClass {
             | KnownClass::UnionType
             | KnownClass::GeneratorType
             | KnownClass::AsyncGeneratorType
+            | KnownClass::CoroutineType
             | KnownClass::NoneType
             | KnownClass::Any
             | KnownClass::StdlibAlias
@@ -2917,7 +2929,11 @@ impl KnownClass {
     /// 2. It's probably more performant.
     const fn is_protocol(self) -> bool {
         match self {
-            Self::SupportsIndex | Self::Iterable | Self::Iterator => true,
+            Self::SupportsIndex
+            | Self::Iterable
+            | Self::Iterator
+            | Self::Awaitable
+            | Self::Generator => true,
 
             Self::Any
             | Self::Bool
@@ -2946,6 +2962,7 @@ impl KnownClass {
             | Self::GenericAlias
             | Self::GeneratorType
             | Self::AsyncGeneratorType
+            | Self::CoroutineType
             | Self::ModuleType
             | Self::FunctionType
             | Self::MethodType
@@ -3011,6 +3028,8 @@ impl KnownClass {
             Self::ExceptionGroup => "ExceptionGroup",
             Self::Staticmethod => "staticmethod",
             Self::Classmethod => "classmethod",
+            Self::Awaitable => "Awaitable",
+            Self::Generator => "Generator",
             Self::Deprecated => "deprecated",
             Self::GenericAlias => "GenericAlias",
             Self::ModuleType => "ModuleType",
@@ -3021,6 +3040,7 @@ impl KnownClass {
             Self::WrapperDescriptorType => "WrapperDescriptorType",
             Self::GeneratorType => "GeneratorType",
             Self::AsyncGeneratorType => "AsyncGeneratorType",
+            Self::CoroutineType => "CoroutineType",
             Self::NamedTuple => "NamedTuple",
             Self::NoneType => "NoneType",
             Self::SpecialForm => "_SpecialForm",
@@ -3281,11 +3301,14 @@ impl KnownClass {
             | Self::MethodType
             | Self::GeneratorType
             | Self::AsyncGeneratorType
+            | Self::CoroutineType
             | Self::MethodWrapperType
             | Self::UnionType
             | Self::WrapperDescriptorType => KnownModule::Types,
             Self::NoneType => KnownModule::Typeshed,
             Self::Any
+            | Self::Awaitable
+            | Self::Generator
             | Self::SpecialForm
             | Self::TypeVar
             | Self::NamedTuple
@@ -3366,12 +3389,15 @@ impl KnownClass {
             | Self::ExceptionGroup
             | Self::Staticmethod
             | Self::Classmethod
+            | Self::Awaitable
+            | Self::Generator
             | Self::Deprecated
             | Self::GenericAlias
             | Self::ModuleType
             | Self::FunctionType
             | Self::GeneratorType
             | Self::AsyncGeneratorType
+            | Self::CoroutineType
             | Self::MethodType
             | Self::MethodWrapperType
             | Self::WrapperDescriptorType
@@ -3443,6 +3469,7 @@ impl KnownClass {
             | Self::WrapperDescriptorType
             | Self::GeneratorType
             | Self::AsyncGeneratorType
+            | Self::CoroutineType
             | Self::SpecialForm
             | Self::ChainMap
             | Self::Counter
@@ -3457,6 +3484,8 @@ impl KnownClass {
             | Self::ExceptionGroup
             | Self::Staticmethod
             | Self::Classmethod
+            | Self::Awaitable
+            | Self::Generator
             | Self::Deprecated
             | Self::TypeVar
             | Self::ParamSpec
@@ -3513,12 +3542,15 @@ impl KnownClass {
             "ExceptionGroup" => Self::ExceptionGroup,
             "staticmethod" => Self::Staticmethod,
             "classmethod" => Self::Classmethod,
+            "Awaitable" => Self::Awaitable,
+            "Generator" => Self::Generator,
             "deprecated" => Self::Deprecated,
             "GenericAlias" => Self::GenericAlias,
             "NoneType" => Self::NoneType,
             "ModuleType" => Self::ModuleType,
             "GeneratorType" => Self::GeneratorType,
             "AsyncGeneratorType" => Self::AsyncGeneratorType,
+            "CoroutineType" => Self::CoroutineType,
             "FunctionType" => Self::FunctionType,
             "MethodType" => Self::MethodType,
             "UnionType" => Self::UnionType,
@@ -3623,6 +3655,7 @@ impl KnownClass {
             | Self::UnionType
             | Self::GeneratorType
             | Self::AsyncGeneratorType
+            | Self::CoroutineType
             | Self::WrapperDescriptorType
             | Self::Field
             | Self::KwOnly
@@ -3642,6 +3675,7 @@ impl KnownClass {
             | Self::Iterable
             | Self::Iterator
             | Self::NewType => matches!(module, KnownModule::Typing | KnownModule::TypingExtensions),
+            Self::Awaitable | Self::Generator => matches!(module, KnownModule::Typing | KnownModule::TypingExtensions | KnownModule::Abc),
             Self::Deprecated => matches!(module, KnownModule::Warnings | KnownModule::TypingExtensions),
 
         }

crates/ty_python_semantic/src/types/function.rs

@@ -341,12 +341,16 @@ impl<'db> OverloadLiteral<'db> {
             GenericContext::from_type_params(db, index, type_params)
         });
 
+        let index = semantic_index(db, scope.file(db));
+        let is_generator = scope.file_scope_id(db).is_generator_function(index);
+
         Signature::from_function(
             db,
             generic_context,
             inherited_generic_context,
             definition,
             function_stmt_node,
+            is_generator,
         )
     }
 

crates/ty_python_semantic/src/types/infer.rs

@@ -119,7 +119,7 @@ use crate::types::{
     CallDunderError, CallableType, ClassLiteral, ClassType, DataclassParams, DynamicType,
     IntersectionBuilder, IntersectionType, KnownClass, KnownInstanceType, LintDiagnosticGuard,
     MemberLookupPolicy, MetaclassCandidate, PEP695TypeAliasType, Parameter, ParameterForm,
-    Parameters, SpecialFormType, SubclassOfType, Truthiness, Type, TypeAliasType,
+    Parameters, Protocol, SpecialFormType, SubclassOfType, Truthiness, Type, TypeAliasType,
     TypeAndQualifiers, TypeIsType, TypeQualifiers, TypeVarBoundOrConstraints, TypeVarInstance,
     TypeVarKind, TypeVarVariance, UnionBuilder, UnionType, binding_type, todo_type,
 };
@@ -6112,8 +6112,29 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             node_index: _,
             value,
         } = await_expression;
-        self.infer_expression(value);
-        todo_type!("generic `typing.Awaitable` type")
+        let value_ty = self.infer_expression(value);
+
+        value_ty
+            .try_call_dunder(self.db(), "__await__", CallArguments::none())
+            .map_or(Type::unknown(), |result| {
+                let generator_ty = result.return_type(self.db());
+
+                if let Type::ProtocolInstance(instance) = generator_ty {
+                    if let Protocol::FromClass(class) = instance.inner {
+                        if class.is_known(self.db(), KnownClass::Generator) {
+                            if let Some(specialization) =
+                                class.class_literal_specialized(self.db(), None).1
+                            {
+                                if let [_, _, return_ty] = specialization.types(self.db()) {
+                                    return *return_ty;
+                                }
+                            }
+                        }
+                    }
+                }
+
+                Type::unknown()
+            })
     }
 
     // Perform narrowing with applicable constraints between the current scope and the enclosing scope.

crates/ty_python_semantic/src/types/signatures.rs

@@ -18,7 +18,7 @@ use smallvec::{SmallVec, smallvec_inline};
 use super::{DynamicType, Type, TypeTransformer, TypeVarVariance, definition_expression_type};
 use crate::semantic_index::definition::Definition;
 use crate::types::generics::{GenericContext, walk_generic_context};
-use crate::types::{TypeMapping, TypeRelation, TypeVarInstance, todo_type};
+use crate::types::{KnownClass, TypeMapping, TypeRelation, TypeVarInstance, todo_type};
 use crate::{Db, FxOrderSet};
 use ruff_python_ast::{self as ast, name::Name};
 
@@ -320,14 +320,18 @@ impl<'db> Signature<'db> {
         inherited_generic_context: Option<GenericContext<'db>>,
         definition: Definition<'db>,
         function_node: &ast::StmtFunctionDef,
+        is_async_generator: bool,
     ) -> Self {
         let parameters =
             Parameters::from_parameters(db, definition, function_node.parameters.as_ref());
         let return_ty = function_node.returns.as_ref().map(|returns| {
-            if function_node.is_async {
-                todo_type!("generic types.CoroutineType")
+            let plain_return_ty = definition_expression_type(db, definition, returns.as_ref());
+
+            if function_node.is_async && !is_async_generator {
+                KnownClass::CoroutineType
+                    .to_specialized_instance(db, [Type::any(), Type::any(), plain_return_ty])
             } else {
-                definition_expression_type(db, definition, returns.as_ref())
+                plain_return_ty
             }
         });
         let legacy_generic_context =