[ty] Support async/await

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/async.md

@@ -0,0 +1,33 @@
+# `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
+```
+
+## Generic `async` functions
+
+```py
+from typing import TypeVar
+
+T = TypeVar("T")
+
+async def persist(x: T) -> T:
+    return x
+
+reveal_type(persist)  # revealed: def persist(x: T) -> CoroutineType[Any, Any, T]
+
+async def f(x: int):
+    result = await persist(x)
+
+    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/with/async.md

@@ -17,5 +17,65 @@ class Manager:
 
 async def test():
     async with Manager() as f:
-        reveal_type(f)  # revealed: @Todo(async `with` statement)
+        reveal_type(f)  # revealed: Target
+```
+
+## `@asynccontextmanager`
+
+```py
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+
+class Session: ...
+
+@asynccontextmanager
+async def connect() -> AsyncGenerator[Session]:
+    yield Session()
+
+# TODO: this should be `() -> _AsyncGeneratorContextManager[Session, None]`
+reveal_type(connect)  # revealed: (...) -> _AsyncGeneratorContextManager[Unknown, None]
+
+async def main():
+    async with connect() as session:
+        # TODO: should be `Session`
+        reveal_type(session)  # revealed: Unknown
+```
+
+## `asyncio.timeout`
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+```py
+import asyncio
+
+async def long_running_task():
+    await asyncio.sleep(5)
+
+async def main():
+    async with asyncio.timeout(1):
+        await long_running_task()
+```
+
+## `asyncio.TaskGroup`
+
+```toml
+[environment]
+python-version = "3.11"
+```
+
+```py
+import asyncio
+
+async def long_running_task():
+    await asyncio.sleep(5)
+
+async def main():
+    async with asyncio.TaskGroup() as tg:
+        # TODO: should be `TaskGroup`
+        reveal_type(tg)  # revealed: Unknown
+
+        tg.create_task(long_running_task())
 ```

crates/ty_python_semantic/src/types.rs

@@ -4790,6 +4790,29 @@ impl<'db> Type<'db> {
         }
     }
 
+    fn resolve_await(self, db: &'db dyn Db) -> Type<'db> {
+        self.try_call_dunder(db, "__await__", CallArguments::none())
+            .map_or(Type::unknown(), |result| {
+                let generator_ty = result.return_type(db);
+
+                if let Type::ProtocolInstance(instance) = generator_ty {
+                    if let Protocol::FromClass(class) = instance.inner {
+                        if class.is_known(db, KnownClass::Generator) {
+                            if let Some(specialization) =
+                                class.class_literal_specialized(db, None).1
+                            {
+                                if let [_, _, return_ty] = specialization.types(db) {
+                                    return *return_ty;
+                                }
+                            }
+                        }
+                    }
+                }
+
+                Type::unknown()
+            })
+    }
+
     /// Given a class literal or non-dynamic SubclassOf type, try calling it (creating an instance)
     /// and return the resulting instance type.
     ///

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

@@ -3169,26 +3169,17 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         let context_expr = with_item.context_expr(self.module());
         let target = with_item.target(self.module());
 
-        let target_ty = if with_item.is_async() {
-            let _context_expr_ty = self.infer_standalone_expression(context_expr);
-            todo_type!("async `with` statement")
-        } else {
-            match with_item.target_kind() {
-                TargetKind::Sequence(unpack_position, unpack) => {
-                    let unpacked = infer_unpack_types(self.db(), unpack);
-                    if unpack_position == UnpackPosition::First {
-                        self.context.extend(unpacked.diagnostics());
-                    }
-                    unpacked.expression_type(target)
-                }
-                TargetKind::Single => {
-                    let context_expr_ty = self.infer_standalone_expression(context_expr);
-                    self.infer_context_expression(
-                        context_expr,
-                        context_expr_ty,
-                        with_item.is_async(),
-                    )
+        let target_ty = match with_item.target_kind() {
+            TargetKind::Sequence(unpack_position, unpack) => {
+                let unpacked = infer_unpack_types(self.db(), unpack);
+                if unpack_position == UnpackPosition::First {
+                    self.context.extend(unpacked.diagnostics());
                 }
+                unpacked.expression_type(target)
+            }
+            TargetKind::Single => {
+                let context_expr_ty = self.infer_standalone_expression(context_expr);
+                self.infer_context_expression(context_expr, context_expr_ty, with_item.is_async())
             }
         };
 
@@ -3208,9 +3199,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         context_expression_type: Type<'db>,
         is_async: bool,
     ) -> Type<'db> {
-        // TODO: Handle async with statements (they use `aenter` and `aexit`)
         if is_async {
-            return todo_type!("async `with` statement");
+            // TODO: proper error handling for `__aenter__`/`__aexit__` calls
+            return context_expression_type
+                .try_call_dunder(self.db(), "__aenter__", CallArguments::none())
+                .map_or(Type::unknown(), |bindings| {
+                    bindings.return_type(self.db()).resolve_await(self.db())
+                });
         }
 
         context_expression_type
@@ -6112,8 +6107,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             node_index: _,
             value,
         } = await_expression;
-        self.infer_expression(value);
-        todo_type!("generic `typing.Awaitable` type")
+        self.infer_expression(value).resolve_await(self.db())
     }
 
     // Perform narrowing with applicable constraints between the current scope and the enclosing scope.