[red-knot] infer attribute assignments bound in comprehensions (astral-sh#17396)

## Summary

This PR is a follow-up to astral-sh#16852.

Instance variables bound in comprehensions are recorded, allowing type
inference to work correctly.

This required adding support for unpacking in comprehension which
resolves astral-sh#15369.

## Test Plan

One TODO in `mdtest/attributes.md` is now resolved, and some new test
cases are added.

---------

Co-authored-by: Dhruv Manilawala <[email protected]>

Author: mtshiba

crates/red_knot_project/tests/check.rs

@@ -6,7 +6,9 @@ use ruff_db::parsed::parsed_module;
 use ruff_db::system::{SystemPath, SystemPathBuf, TestSystem};
 use ruff_python_ast::visitor::source_order;
 use ruff_python_ast::visitor::source_order::SourceOrderVisitor;
-use ruff_python_ast::{self as ast, Alias, Expr, Parameter, ParameterWithDefault, Stmt};
+use ruff_python_ast::{
+    self as ast, Alias, Comprehension, Expr, Parameter, ParameterWithDefault, Stmt,
+};
 
 fn setup_db(project_root: &SystemPath, system: TestSystem) -> anyhow::Result<ProjectDatabase> {
     let project = ProjectMetadata::discover(project_root, &system)?;
@@ -258,6 +260,14 @@ impl SourceOrderVisitor<'_> for PullTypesVisitor<'_> {
         source_order::walk_expr(self, expr);
     }
 
+    fn visit_comprehension(&mut self, comprehension: &Comprehension) {
+        self.visit_expr(&comprehension.iter);
+        self.visit_target(&comprehension.target);
+        for if_expr in &comprehension.ifs {
+            self.visit_expr(if_expr);
+        }
+    }
+
     fn visit_parameter(&mut self, parameter: &Parameter) {
         let _ty = parameter.inferred_type(&self.model);
 

crates/red_knot_python_semantic/resources/mdtest/attributes.md

@@ -397,15 +397,27 @@ class IntIterable:
     def __iter__(self) -> IntIterator:
         return IntIterator()
 
+class TupleIterator:
+    def __next__(self) -> tuple[int, str]:
+        return (1, "a")
+
+class TupleIterable:
+    def __iter__(self) -> TupleIterator:
+        return TupleIterator()
+
 class C:
     def __init__(self) -> None:
         [... for self.a in IntIterable()]
+        [... for (self.b, self.c) in TupleIterable()]
+        [... for self.d in IntIterable() for self.e in IntIterable()]
 
 c_instance = C()
 
-# TODO: Should be `Unknown | int`
-# error: [unresolved-attribute]
-reveal_type(c_instance.a)  # revealed: Unknown
+reveal_type(c_instance.a)  # revealed: Unknown | int
+reveal_type(c_instance.b)  # revealed: Unknown | int
+reveal_type(c_instance.c)  # revealed: Unknown | str
+reveal_type(c_instance.d)  # revealed: Unknown | int
+reveal_type(c_instance.e)  # revealed: Unknown | int
 ```
 
 #### Conditionally declared / bound attributes

crates/red_knot_python_semantic/resources/mdtest/unpacking.md

@@ -708,3 +708,95 @@ with ContextManager() as (a, b, c):
     reveal_type(b)  # revealed: Unknown
     reveal_type(c)  # revealed: Unknown
 ```
+
+## Comprehension
+
+Unpacking in a comprehension.
+
+### Same types
+
+```py
+def _(arg: tuple[tuple[int, int], tuple[int, int]]):
+    # revealed: tuple[int, int]
+    [reveal_type((a, b)) for a, b in arg]
+```
+
+### Mixed types (1)
+
+```py
+def _(arg: tuple[tuple[int, int], tuple[int, str]]):
+    # revealed: tuple[int, int | str]
+    [reveal_type((a, b)) for a, b in arg]
+```
+
+### Mixed types (2)
+
+```py
+def _(arg: tuple[tuple[int, str], tuple[str, int]]):
+    # revealed: tuple[int | str, str | int]
+    [reveal_type((a, b)) for a, b in arg]
+```
+
+### Mixed types (3)
+
+```py
+def _(arg: tuple[tuple[int, int, int], tuple[int, str, bytes], tuple[int, int, str]]):
+    # revealed: tuple[int, int | str, int | bytes | str]
+    [reveal_type((a, b, c)) for a, b, c in arg]
+```
+
+### Same literal values
+
+```py
+# revealed: tuple[Literal[1, 3], Literal[2, 4]]
+[reveal_type((a, b)) for a, b in ((1, 2), (3, 4))]
+```
+
+### Mixed literal values (1)
+
+```py
+# revealed: tuple[Literal[1, "a"], Literal[2, "b"]]
+[reveal_type((a, b)) for a, b in ((1, 2), ("a", "b"))]
+```
+
+### Mixed literals values (2)
+
+```py
+# error: "Object of type `Literal[1]` is not iterable"
+# error: "Object of type `Literal[2]` is not iterable"
+# error: "Object of type `Literal[4]` is not iterable"
+# error: [invalid-assignment] "Not enough values to unpack (expected 2, got 1)"
+# revealed: tuple[Unknown | Literal[3, 5], Unknown | Literal["a", "b"]]
+[reveal_type((a, b)) for a, b in (1, 2, (3, "a"), 4, (5, "b"), "c")]
+```
+
+### Custom iterator (1)
+
+```py
+class Iterator:
+    def __next__(self) -> tuple[int, int]:
+        return (1, 2)
+
+class Iterable:
+    def __iter__(self) -> Iterator:
+        return Iterator()
+
+# revealed: tuple[int, int]
+[reveal_type((a, b)) for a, b in Iterable()]
+```
+
+### Custom iterator (2)
+
+```py
+class Iterator:
+    def __next__(self) -> bytes:
+        return b""
+
+class Iterable:
+    def __iter__(self) -> Iterator:
+        return Iterator()
+
+def _(arg: tuple[tuple[int, str], Iterable]):
+    # revealed: tuple[int | bytes, str | bytes]
+    [reveal_type((a, b)) for a, b in arg]
+```

crates/red_knot_python_semantic/src/semantic_index.rs

@@ -940,7 +940,7 @@ def f(a: str, /, b: str, c: int = 1, *args, d: int = 2, **kwargs):
             panic!("expected generator definition")
         };
         let target = comprehension.target();
-        let name = target.id().as_str();
+        let name = target.as_name_expr().unwrap().id().as_str();
 
         assert_eq!(name, "x");
         assert_eq!(target.range(), TextRange::new(23.into(), 24.into()));

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -18,11 +18,12 @@ use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
 use crate::semantic_index::ast_ids::AstIdsBuilder;
 use crate::semantic_index::definition::{
     AnnotatedAssignmentDefinitionKind, AnnotatedAssignmentDefinitionNodeRef,
-    AssignmentDefinitionKind, AssignmentDefinitionNodeRef, ComprehensionDefinitionNodeRef,
-    Definition, DefinitionCategory, DefinitionKind, DefinitionNodeKey, DefinitionNodeRef,
-    Definitions, ExceptHandlerDefinitionNodeRef, ForStmtDefinitionKind, ForStmtDefinitionNodeRef,
-    ImportDefinitionNodeRef, ImportFromDefinitionNodeRef, MatchPatternDefinitionNodeRef,
-    StarImportDefinitionNodeRef, TargetKind, WithItemDefinitionKind, WithItemDefinitionNodeRef,
+    AssignmentDefinitionKind, AssignmentDefinitionNodeRef, ComprehensionDefinitionKind,
+    ComprehensionDefinitionNodeRef, Definition, DefinitionCategory, DefinitionKind,
+    DefinitionNodeKey, DefinitionNodeRef, Definitions, ExceptHandlerDefinitionNodeRef,
+    ForStmtDefinitionKind, ForStmtDefinitionNodeRef, ImportDefinitionNodeRef,
+    ImportFromDefinitionNodeRef, MatchPatternDefinitionNodeRef, StarImportDefinitionNodeRef,
+    TargetKind, WithItemDefinitionKind, WithItemDefinitionNodeRef,
 };
 use crate::semantic_index::expression::{Expression, ExpressionKind};
 use crate::semantic_index::predicate::{
@@ -850,31 +851,35 @@ impl<'db> SemanticIndexBuilder<'db> {
 
         // The `iter` of the first generator is evaluated in the outer scope, while all subsequent
         // nodes are evaluated in the inner scope.
-        self.add_standalone_expression(&generator.iter);
+        let value = self.add_standalone_expression(&generator.iter);
         self.visit_expr(&generator.iter);
         self.push_scope(scope);
 
-        self.push_assignment(CurrentAssignment::Comprehension {
-            node: generator,
-            first: true,
-        });
-        self.visit_expr(&generator.target);
-        self.pop_assignment();
+        self.add_unpackable_assignment(
+            &Unpackable::Comprehension {
+                node: generator,
+                first: true,
+            },
+            &generator.target,
+            value,
+        );
 
         for expr in &generator.ifs {
             self.visit_expr(expr);
         }
 
         for generator in generators_iter {
-            self.add_standalone_expression(&generator.iter);
+            let value = self.add_standalone_expression(&generator.iter);
             self.visit_expr(&generator.iter);
 
-            self.push_assignment(CurrentAssignment::Comprehension {
-                node: generator,
-                first: false,
-            });
-            self.visit_expr(&generator.target);
-            self.pop_assignment();
+            self.add_unpackable_assignment(
+                &Unpackable::Comprehension {
+                    node: generator,
+                    first: false,
+                },
+                &generator.target,
+                value,
+            );
 
             for expr in &generator.ifs {
                 self.visit_expr(expr);
@@ -933,9 +938,30 @@ impl<'db> SemanticIndexBuilder<'db> {
 
         let current_assignment = match target {
             ast::Expr::List(_) | ast::Expr::Tuple(_) => {
+                if matches!(unpackable, Unpackable::Comprehension { .. }) {
+                    debug_assert_eq!(
+                        self.scopes[self.current_scope()].node().scope_kind(),
+                        ScopeKind::Comprehension
+                    );
+                }
+                // The first iterator of the comprehension is evaluated in the outer scope, while all subsequent
+                // nodes are evaluated in the inner scope.
+                // SAFETY: The current scope is the comprehension, and the comprehension scope must have a parent scope.
+                let value_file_scope =
+                    if let Unpackable::Comprehension { first: true, .. } = unpackable {
+                        self.scope_stack
+                            .iter()
+                            .rev()
+                            .nth(1)
+                            .expect("The comprehension scope must have a parent scope")
+                            .file_scope_id
+                    } else {
+                        self.current_scope()
+                    };
                 let unpack = Some(Unpack::new(
                     self.db,
                     self.file,
+                    value_file_scope,
                     self.current_scope(),
                     // SAFETY: `target` belongs to the `self.module` tree
                     #[allow(unsafe_code)]
@@ -1804,7 +1830,7 @@ where
         let node_key = NodeKey::from_node(expr);
 
         match expr {
-            ast::Expr::Name(name_node @ ast::ExprName { id, ctx, .. }) => {
+            ast::Expr::Name(ast::ExprName { id, ctx, .. }) => {
                 let (is_use, is_definition) = match (ctx, self.current_assignment()) {
                     (ast::ExprContext::Store, Some(CurrentAssignment::AugAssign(_))) => {
                         // For augmented assignment, the target expression is also used.
@@ -1867,12 +1893,17 @@ where
                             // implemented.
                             self.add_definition(symbol, named);
                         }
-                        Some(CurrentAssignment::Comprehension { node, first }) => {
+                        Some(CurrentAssignment::Comprehension {
+                            unpack,
+                            node,
+                            first,
+                        }) => {
                             self.add_definition(
                                 symbol,
                                 ComprehensionDefinitionNodeRef {
+                                    unpack,
                                     iterable: &node.iter,
-                                    target: name_node,
+                                    target: expr,
                                     first,
                                     is_async: node.is_async,
                                 },
@@ -2143,14 +2174,37 @@ where
                                 DefinitionKind::WithItem(assignment),
                             );
                         }
-                        Some(CurrentAssignment::Comprehension { .. }) => {
-                            // TODO:
+                        Some(CurrentAssignment::Comprehension {
+                            unpack,
+                            node,
+                            first,
+                        }) => {
+                            // SAFETY: `iter` and `expr` belong to the `self.module` tree
+                            #[allow(unsafe_code)]
+                            let assignment = ComprehensionDefinitionKind {
+                                target_kind: TargetKind::from(unpack),
+                                iterable: unsafe {
+                                    AstNodeRef::new(self.module.clone(), &node.iter)
+                                },
+                                target: unsafe { AstNodeRef::new(self.module.clone(), expr) },
+                                first,
+                                is_async: node.is_async,
+                            };
+                            // Temporarily move to the scope of the method to which the instance attribute is defined.
+                            // SAFETY: `self.scope_stack` is not empty because the targets in comprehensions should always introduce a new scope.
+                            let scope = self.scope_stack.pop().expect("The popped scope must be a comprehension, which must have a parent scope");
+                            self.register_attribute_assignment(
+                                object,
+                                attr,
+                                DefinitionKind::Comprehension(assignment),
+                            );
+                            self.scope_stack.push(scope);
                         }
                         Some(CurrentAssignment::AugAssign(_)) => {
                             // TODO:
                         }
                         Some(CurrentAssignment::Named(_)) => {
-                            // TODO:
+                            // A named expression whose target is an attribute is syntactically prohibited
                         }
                         None => {}
                     }
@@ -2244,6 +2298,7 @@ enum CurrentAssignment<'a> {
     Comprehension {
         node: &'a ast::Comprehension,
         first: bool,
+        unpack: Option<(UnpackPosition, Unpack<'a>)>,
     },
     WithItem {
         item: &'a ast::WithItem,
@@ -2257,11 +2312,9 @@ impl CurrentAssignment<'_> {
         match self {
             Self::Assign { unpack, .. }
             | Self::For { unpack, .. }
-            | Self::WithItem { unpack, .. } => unpack.as_mut().map(|(position, _)| position),
-            Self::AnnAssign(_)
-            | Self::AugAssign(_)
-            | Self::Named(_)
-            | Self::Comprehension { .. } => None,
+            | Self::WithItem { unpack, .. }
+            | Self::Comprehension { unpack, .. } => unpack.as_mut().map(|(position, _)| position),
+            Self::AnnAssign(_) | Self::AugAssign(_) | Self::Named(_) => None,
         }
     }
 }
@@ -2316,13 +2369,17 @@ enum Unpackable<'a> {
         item: &'a ast::WithItem,
         is_async: bool,
     },
+    Comprehension {
+        first: bool,
+        node: &'a ast::Comprehension,
+    },
 }
 
 impl<'a> Unpackable<'a> {
     const fn kind(&self) -> UnpackKind {
         match self {
             Unpackable::Assign(_) => UnpackKind::Assign,
-            Unpackable::For(_) => UnpackKind::Iterable,
+            Unpackable::For(_) | Unpackable::Comprehension { .. } => UnpackKind::Iterable,
             Unpackable::WithItem { .. } => UnpackKind::ContextManager,
         }
     }
@@ -2337,6 +2394,11 @@ impl<'a> Unpackable<'a> {
                 is_async: *is_async,
                 unpack,
             },
+            Unpackable::Comprehension { node, first } => CurrentAssignment::Comprehension {
+                node,
+                first: *first,
+                unpack,
+            },
         }
     }
 }