Add coroutine type

typechecker/src/type_evaluator.rs

@@ -21,7 +21,7 @@ use crate::{
     symbol_table::{
         self, Class, Declaration, Id, LookupSymbolRequest, SymbolTable, SymbolTableNode,
     },
-    types::{ClassType, ModuleRef, TypeVar},
+    types::{self, ClassType, ModuleRef, TypeVar},
 };
 
 const LITERAL_TYPE_PARAMETER_MSG: &str = "Type arguments for 'Literal' must be None, a literal value (int, bool, str, or bytes), or an enum value";
@@ -108,7 +108,13 @@ impl<'a> TypeEvaluator<'a> {
                     _ => {
                         let f_type = self.get_type(called_function, Some(symbol_table), None)?;
                         if let PythonType::Callable(c) = &f_type {
-                            let return_type = self.get_return_type_of_callable(c, &call.args);
+                            // TODO(coroutine_annotation): When the function is called and it's
+                            // async then it will return the coroutine type. Not the return type.
+                            let return_type = if c.is_async {
+                                self.get_return_type_of_callable(c, &call.args)
+                            } else {
+                                c.return_type.clone()
+                            };
                             Ok(return_type)
                         } else if let PythonType::Class(c) = &f_type {
                             Ok(f_type)
@@ -365,11 +371,8 @@ impl<'a> TypeEvaluator<'a> {
                 let awaited_type =
                     self.get_type(&a.value, Some(symbol_table), symbol_table_scope)?;
                 let typ = match awaited_type {
-                    PythonType::Callable(callable) => {
-                        let ret_type = self.get_return_type_of_callable(&callable, &[]);
-                        ret_type
-                    }
-                    _ => PythonType::Unknown,
+                    PythonType::Coroutine(callable) => callable.return_type.clone(),
+                    _ => unimplemented!("Can other things be awaited?"),
                 };
 
                 Ok(typ)
@@ -1151,15 +1154,21 @@ impl<'a> TypeEvaluator<'a> {
             }
         }
         let name = f.function_node.name.clone();
+        let return_type = f
+            .function_node
+            .returns
+            .clone()
+            .map_or(PythonType::Unknown, |type_annotation| {
+                self.get_annotation_type(&type_annotation, symbol_table, None)
+            });
         PythonType::Callable(Box::new(CallableType::new(
             name,
             signature,
-            f.function_node
-                .returns
-                .clone()
-                .map_or(PythonType::Unknown, |type_annotation| {
-                    self.get_annotation_type(&type_annotation, symbol_table, None)
-                }),
+            PythonType::Coroutine(Box::new(types::CoroutineType {
+                return_type,
+                send_type: PythonType::Any,
+                yield_type: PythonType::Any,
+            })),
             true,
         )))
     }

typechecker/src/types.rs

@@ -27,6 +27,7 @@ pub enum PythonType {
     /// Union type
     MultiValue(Vec<PythonType>),
     Callable(Box<CallableType>),
+    Coroutine(Box<CoroutineType>),
     Class(ClassType),
     Optional(Box<PythonType>),
     Never,
@@ -122,6 +123,15 @@ impl PartialEq for CallableType {
     }
 }
 
+#[derive(Debug, Eq, PartialEq, Clone)]
+pub struct CoroutineType {
+    pub return_type: PythonType,
+    // TODO(coroutine_annotation): Any Any are send and yield type that are not implemented yet
+    // https://github.com/python/typing/issues/251
+    pub send_type: PythonType,
+    pub yield_type: PythonType,
+}
+
 #[allow(unused)]
 #[derive(Debug, Clone)]
 pub struct ClassType {
@@ -235,14 +245,16 @@ impl Display for PythonType {
                     .map(|arg| arg.to_string())
                     .collect::<Vec<String>>()
                     .join(", ");
-                let return_type_str = if callable_type.is_async {
-                    format!("Coroutine[Any, Any, {}]", callable_type.return_type)
-                } else {
-                    callable_type.return_type.to_string()
-                };
                 let fmt = format!(
                     "(function) Callable[[{}], {}]",
-                    signature_str, return_type_str
+                    signature_str, callable_type.return_type
+                );
+                return write!(f, "{}", fmt);
+            }
+            PythonType::Coroutine(callable_type) => {
+                let fmt = format!(
+                    "Coroutine[{}, {}, {}]",
+                    callable_type.send_type, callable_type.yield_type, callable_type.return_type
                 );
                 return write!(f, "{}", fmt);
             }

typechecker/test_data/inputs/conformance_tests/annotations_coroutine.py

@@ -18,7 +18,6 @@ async def func1(ignored: int, /) -> str:
 
 assert_type(func1, Callable[[int], Coroutine[Any, Any, str]])
 
-
 async def func2() -> None:
     x = await func1(42)
     assert_type(x, str)

typechecker/test_data/output/enderpy_python_type_checker__checker__tests__annotations_coroutine.snap

@@ -1,6 +1,6 @@
 ---
 source: typechecker/src/checker.rs
-description: "1: \"\"\"\n2: Tests for annotating coroutines.\n3: \"\"\"\n4: \n5: # Specification: https://typing.readthedocs.io/en/latest/spec/annotations.html#annotating-generator-functions-and-coroutines\n6: \n7: # > Coroutines introduced in PEP 492 are annotated with the same syntax as\n8: # > ordinary functions. However, the return type annotation corresponds to\n9: # > the type of await expression, not to the coroutine type.\n10: \n11: \n12: from typing import Any, Callable, Coroutine, assert_type\n13: \n14: \n15: async def func1(ignored: int, /) -> str:\n16:     return \"spam\"\n17: \n18: \n19: assert_type(func1, Callable[[int], Coroutine[Any, Any, str]])\n20: \n21: \n22: async def func2() -> None:\n23:     x = await func1(42)\n24:     assert_type(x, str)\n"
+description: "1: \"\"\"\n2: Tests for annotating coroutines.\n3: \"\"\"\n4: \n5: # Specification: https://typing.readthedocs.io/en/latest/spec/annotations.html#annotating-generator-functions-and-coroutines\n6: \n7: # > Coroutines introduced in PEP 492 are annotated with the same syntax as\n8: # > ordinary functions. However, the return type annotation corresponds to\n9: # > the type of await expression, not to the coroutine type.\n10: \n11: \n12: from typing import Any, Callable, Coroutine, assert_type\n13: \n14: \n15: async def func1(ignored: int, /) -> str:\n16:     return \"spam\"\n17: \n18: \n19: assert_type(func1, Callable[[int], Coroutine[Any, Any, str]])\n20: \n21: async def func2() -> None:\n22:     x = await func1(42)\n23:     assert_type(x, str)\n"
 expression: result
 ---
 Line 1: """
@@ -54,28 +54,28 @@ Expr types in the line --->:
         str => (class) str
 
 ---
-Line 20: async def func2() -> None:
+Line 19: async def func2() -> None:
 
 Expr types in the line --->:
          func => (function) Callable[[], Coroutine[Any, Any, None]]
 
 ---
-Line 21:     x = await func1(42)
+Line 20:     x = await func1(42)
 
 Expr types in the line --->:
-        x => Unknown
-        await func1(42) => Unknown
+        x => (class) str
+        await func1(42) => (class) str
         func1 => (function) Callable[[(class) int], Coroutine[Any, Any, (class) str]]
-        func1(42) => (class) str
+        func1(42) => Coroutine[Any, Any, (class) str]
         42 => (class) int
 
 ---
-Line 22:     assert_type(x, str)
+Line 21:     assert_type(x, str)
 
 Expr types in the line --->:
         assert_type => (function) Callable[[TypeVar[_T, ], (class) object], TypeVar[_T, ]]
         assert_type(x, str) => TypeVar[_T, ]
-        x => Unknown
+        x => (class) str
         str => (class) str
 
 ---