[red-knot] Infer generator expressions

crates/red_knot_python_semantic/src/types.rs

@@ -188,6 +188,8 @@ pub enum Type<'db> {
     LiteralString,
     /// A bytes literal
     BytesLiteral(BytesLiteralType<'db>),
+    /// A generator
+    Generator(GeneratorType<'db>),
     // TODO protocols, callable types, overloads, generics, type vars
 }
 
@@ -297,6 +299,7 @@ impl<'db> Type<'db> {
                 // TODO defer to Type::Instance(<bytes from typeshed>).member
                 Type::Unknown
             }
+            Type::Generator(_) => Type::Unknown,
         }
     }
 
@@ -405,6 +408,13 @@ pub struct BytesLiteralType<'db> {
     value: Box<[u8]>,
 }
 
+#[salsa::interned]
+pub struct GeneratorType<'db> {
+    yield_type: Type<'db>,
+    send_type: Type<'db>,
+    return_type: Type<'db>,
+}
+
 #[cfg(test)]
 mod tests {
     use anyhow::Context;

crates/red_knot_python_semantic/src/types/display.rs

@@ -55,6 +55,22 @@ impl Display for DisplayType<'_> {
                 escape.bytes_repr().write(f)?;
                 f.write_str("]")
             }
+            Type::Generator(generator) => write!(
+                f,
+                "Generator[{}, {}, {}]",
+                DisplayType {
+                    ty: &generator.yield_type(self.db),
+                    db: self.db
+                },
+                DisplayType {
+                    ty: &generator.send_type(self.db),
+                    db: self.db
+                },
+                DisplayType {
+                    ty: &generator.return_type(self.db),
+                    db: self.db
+                }
+            ),
         }
     }
 }

crates/red_knot_python_semantic/src/types/infer.rs

@@ -48,6 +48,8 @@ use crate::types::{
 };
 use crate::Db;
 
+use super::GeneratorType;
+
 /// Infer all types for a [`ScopeId`], including all definitions and expressions in that scope.
 /// Use when checking a scope, or needing to provide a type for an arbitrary expression in the
 /// scope.
@@ -1395,12 +1397,15 @@ impl<'db> TypeInferenceBuilder<'db> {
     }
 
     /// Infer the type of the `iter` expression of the first comprehension.
-    fn infer_first_comprehension_iter(&mut self, comprehensions: &[ast::Comprehension]) {
+    fn infer_first_comprehension_iter(
+        &mut self,
+        comprehensions: &[ast::Comprehension],
+    ) -> Type<'db> {
         let mut generators_iter = comprehensions.iter();
         let Some(first_generator) = generators_iter.next() else {
             unreachable!("Comprehension must contain at least one generator");
         };
-        self.infer_expression(&first_generator.iter);
+        self.infer_expression(&first_generator.iter)
     }
 
     fn infer_generator_expression(&mut self, generator: &ast::ExprGenerator) -> Type<'db> {
@@ -1411,10 +1416,14 @@ impl<'db> TypeInferenceBuilder<'db> {
             parenthesized: _,
         } = generator;
 
-        self.infer_first_comprehension_iter(generators);
+        let yield_type = self.infer_first_comprehension_iter(generators);
 
-        // TODO generator type
-        Type::Unknown
+        Type::Generator(GeneratorType::new(
+            self.db,
+            yield_type,
+            Type::None,
+            Type::None,
+        ))
     }
 
     fn infer_list_comprehension_expression(&mut self, listcomp: &ast::ExprListComp) -> Type<'db> {
@@ -3262,6 +3271,18 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn generator_expr() -> anyhow::Result<()> {
+        let mut db = setup_db();
+
+        db.write_file("src/a.py", "x = (y for y in [1,2,3])")?;
+
+        assert_public_ty(&db, "src/a.py", "x", "Generator[list, None, None]");
+        assert_public_ty(&db, "src/a.py", "y", "Unbound");
+
+        Ok(())
+    }
+
     #[test]
     fn dependency_public_symbol_type_change() -> anyhow::Result<()> {
         let mut db = setup_db();