[red-knot] Unpacking and for loop assignments to attributes (#16004)

## Summary

* Support assignments to attributes in more cases:
    - assignments in `for` loops
    - in unpacking assignments
* Add test for multi-target assignments
* Add tests for all other possible assignments to attributes that could
   possibly occur (in decreasing order of likeliness):
    - augmented attribute assignments
    - attribute assignments in `with` statements
    - attribute assignments in comprehensions
- Note: assignments to attributes in named expressions are not
   syntactically allowed

closes #15962

## Test Plan

New Markdown tests

crates/red_knot_python_semantic/resources/mdtest/attributes.md

@@ -232,6 +232,36 @@ reveal_type(c_instance.y)  # revealed: int
 reveal_type(c_instance.z)  # revealed: int
 ```
 
+#### Attributes defined in multi-target assignments
+
+```py
+class C:
+    def __init__(self) -> None:
+        self.a = self.b = 1
+
+c_instance = C()
+
+reveal_type(c_instance.a)  # revealed: Unknown | Literal[1]
+reveal_type(c_instance.b)  # revealed: Unknown | Literal[1]
+```
+
+#### Augmented assignments
+
+```py
+class Weird:
+    def __iadd__(self, other: None) -> str:
+        return "a"
+
+class C:
+    def __init__(self) -> None:
+        self.w = Weird()
+        self.w += None
+
+# TODO: Mypy and pyright do not support this, but it would be great if we could
+# infer `Unknown | str` or at least `Unknown | Weird | str` here.
+reveal_type(C().w)  # revealed: Unknown | Weird
+```
+
 #### Attributes defined in tuple unpackings
 
 ```py
@@ -253,19 +283,24 @@ reveal_type(c_instance.b1)  # revealed: Unknown | Literal["a"]
 reveal_type(c_instance.c1)  # revealed: Unknown | int
 reveal_type(c_instance.d1)  # revealed: Unknown | str
 
-# TODO: This should be supported (no error; type should be: `Unknown | Literal[1]`)
+reveal_type(c_instance.a2)  # revealed: Unknown | Literal[1]
-# error: [unresolved-attribute]
-reveal_type(c_instance.a2)  # revealed: Unknown
 
-# TODO: This should be supported (no error; type should be: `Unknown | Literal["a"]`)
+reveal_type(c_instance.b2)  # revealed: Unknown | Literal["a"]
-# error: [unresolved-attribute]
-reveal_type(c_instance.b2)  # revealed: Unknown
 
-# TODO: Similar for these two (should be `Unknown | int` and `Unknown | str`, respectively)
+reveal_type(c_instance.c2)  # revealed: Unknown | int
-# error: [unresolved-attribute]
+reveal_type(c_instance.d2)  # revealed: Unknown | str
-reveal_type(c_instance.c2)  # revealed: Unknown
+```
-# error: [unresolved-attribute]
+
-reveal_type(c_instance.d2)  # revealed: Unknown
+#### Starred assignments
+
+```py
+class C:
+    def __init__(self) -> None:
+        self.a, *self.b = (1, 2, 3)
+
+c_instance = C()
+reveal_type(c_instance.a)  # revealed: Unknown | Literal[1]
+reveal_type(c_instance.b)  # revealed: Unknown | @Todo(starred unpacking)
 ```
 
 #### Attributes defined in for-loop (unpacking)
@@ -287,6 +322,8 @@ class TupleIterable:
     def __iter__(self) -> TupleIterator:
         return TupleIterator()
 
+class NonIterable: ...
+
 class C:
     def __init__(self):
         for self.x in IntIterable():
@@ -295,14 +332,54 @@ class C:
         for _, self.y in TupleIterable():
             pass
 
-# TODO: Pyright fully supports these, mypy detects the presence of the attributes,
+        # TODO: We should emit a diagnostic here
-# but infers type `Any` for both of them. We should infer `int` and `str` here:
+        for self.z in NonIterable():
+            pass
 
-# error: [unresolved-attribute]
+reveal_type(C().x)  # revealed: Unknown | int
-reveal_type(C().x)  # revealed: Unknown
 
+reveal_type(C().y)  # revealed: Unknown | str
+```
+
+#### Attributes defined in `with` statements
+
+```py
+class ContextManager:
+    def __enter__(self) -> int | None: ...
+    def __exit__(self, exc_type, exc_value, traceback) -> None: ...
+
+class C:
+    def __init__(self) -> None:
+        with ContextManager() as self.x:
+            pass
+
+c_instance = C()
+
+# TODO: Should be `Unknown | int | None`
 # error: [unresolved-attribute]
-reveal_type(C().y)  # revealed: Unknown
+reveal_type(c_instance.x)  # revealed: Unknown
+```
+
+#### Attributes defined in comprehensions
+
+```py
+class IntIterator:
+    def __next__(self) -> int:
+        return 1
+
+class IntIterable:
+    def __iter__(self) -> IntIterator:
+        return IntIterator()
+
+class C:
+    def __init__(self) -> None:
+        [... for self.a in IntIterable()]
+
+c_instance = C()
+
+# TODO: Should be `Unknown | int`
+# error: [unresolved-attribute]
+reveal_type(c_instance.a)  # revealed: Unknown
 ```
 
 #### Conditionally declared / bound attributes
@@ -443,6 +520,15 @@ class C:
 reveal_type(C().x)  # revealed: str
 ```
 
+#### Diagnostics are reported for the right-hand side of attribute assignments
+
+```py
+class C:
+    def __init__(self) -> None:
+        # error: [too-many-positional-arguments]
+        self.x: int = len(1, 2, 3)
+```
+
 ### Pure class variables (`ClassVar`)
 
 #### Annotated with `ClassVar` type qualifier

crates/red_knot_python_semantic/src/semantic_index/attribute_assignment.rs

@@ -1,4 +1,7 @@
-use crate::semantic_index::expression::Expression;
+use crate::{
+    semantic_index::{ast_ids::ScopedExpressionId, expression::Expression},
+    unpack::Unpack,
+};
 
 use ruff_python_ast::name::Name;
 
@@ -14,6 +17,17 @@ pub(crate) enum AttributeAssignment<'db> {
 
     /// An attribute assignment without a type annotation, e.g. `self.x = <value>`.
     Unannotated { value: Expression<'db> },
+
+    /// An attribute assignment where the right-hand side is an iterable, for example
+    /// `for self.x in <iterable>`.
+    Iterable { iterable: Expression<'db> },
+
+    /// An attribute assignment where the left-hand side is an unpacking expression,
+    /// e.g. `self.x, self.y = <value>`.
+    Unpack {
+        attribute_expression_id: ScopedExpressionId,
+        unpack: Unpack<'db>,
+    },
 }
 
 pub(crate) type AttributeAssignments<'db> = FxHashMap<Name, Vec<AttributeAssignment<'db>>>;

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -6,9 +6,9 @@ use rustc_hash::{FxHashMap, FxHashSet};
 use ruff_db::files::File;
 use ruff_db::parsed::ParsedModule;
 use ruff_index::IndexVec;
-use ruff_python_ast as ast;
 use ruff_python_ast::name::Name;
 use ruff_python_ast::visitor::{walk_expr, walk_pattern, walk_stmt, Visitor};
+use ruff_python_ast::{self as ast, ExprContext};
 
 use crate::ast_node_ref::AstNodeRef;
 use crate::module_name::ModuleName;
@@ -1231,6 +1231,20 @@ where
                         unpack: None,
                         first: false,
                     }),
+                    ast::Expr::Attribute(ast::ExprAttribute {
+                        value: object,
+                        attr,
+                        ..
+                    }) => {
+                        self.register_attribute_assignment(
+                            object,
+                            attr,
+                            AttributeAssignment::Iterable {
+                                iterable: iter_expr,
+                            },
+                        );
+                        None
+                    }
                     _ => None,
                 };
 
@@ -1459,7 +1473,7 @@ where
     fn visit_expr(&mut self, expr: &'ast ast::Expr) {
         self.scopes_by_expression
             .insert(expr.into(), self.current_scope());
-        self.current_ast_ids().record_expression(expr);
+        let expression_id = self.current_ast_ids().record_expression(expr);
 
         match expr {
             ast::Expr::Name(name_node @ ast::ExprName { id, ctx, .. }) => {
@@ -1718,6 +1732,35 @@ where
 
                 self.simplify_visibility_constraints(pre_op);
             }
+            ast::Expr::Attribute(ast::ExprAttribute {
+                value: object,
+                attr,
+                ctx: ExprContext::Store,
+                range: _,
+            }) => {
+                if let Some(
+                    CurrentAssignment::Assign {
+                        unpack: Some(unpack),
+                        ..
+                    }
+                    | CurrentAssignment::For {
+                        unpack: Some(unpack),
+                        ..
+                    },
+                ) = self.current_assignment()
+                {
+                    self.register_attribute_assignment(
+                        object,
+                        attr,
+                        AttributeAssignment::Unpack {
+                            attribute_expression_id: expression_id,
+                            unpack,
+                        },
+                    );
+                }
+
+                walk_expr(self, expr);
+            }
             _ => {
                 walk_expr(self, expr);
             }

crates/red_knot_python_semantic/src/types.rs

@@ -40,6 +40,7 @@ use crate::types::call::{
 };
 use crate::types::class_base::ClassBase;
 use crate::types::diagnostic::INVALID_TYPE_FORM;
+use crate::types::infer::infer_unpack_types;
 use crate::types::mro::{Mro, MroError, MroIterator};
 use crate::types::narrow::narrowing_constraint;
 use crate::{Db, FxOrderSet, Module, Program, PythonVersion};
@@ -4231,6 +4232,33 @@ impl<'db> Class<'db> {
 
                     union_of_inferred_types = union_of_inferred_types.add(inferred_ty);
                 }
+                AttributeAssignment::Iterable { iterable } => {
+                    // We found an attribute assignment like:
+                    //
+                    //     for self.name in <iterable>:
+
+                    // TODO: Potential diagnostics resulting from the iterable are currently not reported.
+
+                    let iterable_ty = infer_expression_type(db, *iterable);
+                    let inferred_ty = iterable_ty.iterate(db).unwrap_without_diagnostic();
+
+                    union_of_inferred_types = union_of_inferred_types.add(inferred_ty);
+                }
+                AttributeAssignment::Unpack {
+                    attribute_expression_id,
+                    unpack,
+                } => {
+                    // We found an unpacking assignment like:
+                    //
+                    //     .., self.name, .. = <value>
+                    //     (.., self.name, ..) = <value>
+                    //     [.., self.name, ..] = <value>
+
+                    let inferred_ty = infer_unpack_types(db, *unpack)
+                        .get(*attribute_expression_id)
+                        .expect("Failed to look up type of attribute in unpack assignment");
+                    union_of_inferred_types = union_of_inferred_types.add(inferred_ty);
+                }
             }
         }
 

crates/red_knot_python_semantic/src/types/infer.rs

@@ -200,7 +200,7 @@ pub(crate) fn infer_expression_types<'db>(
 /// type of the variables involved in this unpacking along with any violations that are detected
 /// during this unpacking.
 #[salsa::tracked(return_ref)]
-fn infer_unpack_types<'db>(db: &'db dyn Db, unpack: Unpack<'db>) -> UnpackResult<'db> {
+pub(super) fn infer_unpack_types<'db>(db: &'db dyn Db, unpack: Unpack<'db>) -> UnpackResult<'db> {
     let file = unpack.file(db);
     let _span =
         tracing::trace_span!("infer_unpack_types", range=?unpack.range(db), file=%file.path(db))

crates/red_knot_python_semantic/src/types/unpacker.rs

@@ -72,11 +72,9 @@ impl<'db> Unpacker<'db> {
         value_ty: Type<'db>,
     ) {
         match target {
-            ast::Expr::Name(target_name) => {
+            ast::Expr::Name(_) | ast::Expr::Attribute(_) => {
-                self.targets.insert(
+                self.targets
-                    target_name.scoped_expression_id(self.db(), self.scope),
+                    .insert(target.scoped_expression_id(self.db(), self.scope), value_ty);
-                    value_ty,
-                );
             }
             ast::Expr::Starred(ast::ExprStarred { value, .. }) => {
                 self.unpack_inner(value, value_expr, value_ty);