Also check assignment target in unpack assignments and for loops

Author: sharkdp

crates/red_knot_python_semantic/resources/mdtest/attributes.md

@@ -439,17 +439,32 @@ reveal_type(Foo.__class__)  # revealed: Literal[type]
 ## Module attributes
 
 ```py path=mod.py
-global_symbol: int = 1
+global_symbol: str = "a"
 ```
 
 ```py
 import mod
 
-reveal_type(mod.global_symbol)  # revealed: int
-mod.global_symbol = 2
+reveal_type(mod.global_symbol)  # revealed: str
+mod.global_symbol = "b"
 
-# error: [invalid-assignment] "Object of type `Literal["1"]` is not assignable to `int`"
-mod.global_symbol = "1"
+# error: [invalid-assignment] "Object of type `Literal[1]` is not assignable to `str`"
+mod.global_symbol = 1
+
+# error: [invalid-assignment] "Object of type `Literal[1]` is not assignable to `str`"
+(_, mod.global_symbol) = (..., 1)
+
+class IntIterator:
+    def __next__(self) -> int:
+        return 42
+
+class IntIterable:
+    def __iter__(self) -> IntIterator:
+        return IntIterator()
+
+# error: [invalid-assignment] "Object of type `int` is not assignable to `str`"
+for mod.global_symbol in IntIterable():
+    pass
 ```
 
 ## Literal types

crates/red_knot_python_semantic/src/types.rs

@@ -3319,6 +3319,14 @@ impl<'db> IterationOutcome<'db> {
             }
         }
     }
+
+    fn unwrap_without_diagnostic(self) -> Type<'db> {
+        match self {
+            Self::Iterable { element_ty } => element_ty,
+            Self::NotIterable { .. } => Type::unknown(),
+            Self::PossiblyUnboundDunderIter { element_ty, .. } => element_ty,
+        }
+    }
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1973,24 +1973,46 @@ impl<'db> TypeInferenceBuilder<'db> {
         }
     }
 
-    /// Infer the definition types involved in a `target` expression.
+    /// Infer the (definition) types involved in a `target` expression.
     ///
     /// This is used for assignment statements, for statements, etc. with a single or multiple
-    /// targets (unpacking).
+    /// targets (unpacking). If `target` is an attribute expression, we check that the assignment
+    /// is valid. For 'target's that are definitions, this check happens elsewhere.
     ///
     /// # Panics
     ///
     /// If the `value` is not a standalone expression.
     fn infer_target(&mut self, target: &ast::Expr, value: &ast::Expr) {
+        let check_assignment_ty = match target {
+            ast::Expr::Name(_) => None,
+            _ => Some(self.infer_standalone_expression(value)),
+        };
+        self.infer_target_and_check_assignment(target, check_assignment_ty);
+    }
+
+    /// Similar to [`infer_target`], but with a custom type to check the assignment against.
+    fn infer_target_and_check_assignment(
+        &mut self,
+        target: &ast::Expr,
+        check_assignment_ty: Option<Type<'db>>,
+    ) {
         match target {
             ast::Expr::Name(name) => self.infer_definition(name),
             ast::Expr::List(ast::ExprList { elts, .. })
             | ast::Expr::Tuple(ast::ExprTuple { elts, .. }) => {
+                let mut check_assignment_tys: Box<dyn Iterator<Item = Type>> =
+                    match check_assignment_ty {
+                        Some(Type::Tuple(tuple)) => {
+                            Box::new(tuple.elements(self.db()).into_iter().copied())
+                        }
+                        Some(ty) => Box::new(std::iter::repeat(
+                            ty.iterate(self.db()).unwrap_without_diagnostic(),
+                        )),
+                        None => Box::new(std::iter::empty()),
+                    };
+
                 for element in elts {
-                    self.infer_target(element, value);
-                }
-                if elts.is_empty() {
-                    self.infer_standalone_expression(value);
+                    self.infer_target_and_check_assignment(element, check_assignment_tys.next());
                 }
             }
             ast::Expr::Attribute(
@@ -2002,15 +2024,14 @@ impl<'db> TypeInferenceBuilder<'db> {
                 let lhs_ty = self.infer_attribute_expression(lhs_expr);
                 self.store_expression_type(target, lhs_ty);
 
-                let rhs_ty = self.infer_standalone_expression(value);
-
-                if !rhs_ty.is_assignable_to(self.db(), lhs_ty) {
-                    report_invalid_assignment(&self.context, target.into(), lhs_ty, rhs_ty);
+                if let Some(rhs_ty) = check_assignment_ty {
+                    if !rhs_ty.is_assignable_to(self.db(), lhs_ty) {
+                        report_invalid_assignment(&self.context, target.into(), lhs_ty, rhs_ty);
+                    }
                 }
             }
             _ => {
                 // TODO: Remove this once we handle all possible assignment targets.
-                self.infer_standalone_expression(value);
                 self.infer_expression(target);
             }
         }
@@ -2279,7 +2300,16 @@ impl<'db> TypeInferenceBuilder<'db> {
             is_async: _,
         } = for_statement;
 
-        self.infer_target(target, iter);
+        let check_assignment_ty = match target.as_ref() {
+            ast::Expr::Name(_) => None, // Check will be done in `infer_definition`
+            _ => {
+                let iterable_ty = self.infer_standalone_expression(iter);
+                let iterator_ty = iterable_ty.iterate(self.db()).unwrap_without_diagnostic();
+                Some(iterator_ty)
+            }
+        };
+        self.infer_target_and_check_assignment(target, check_assignment_ty);
+
         self.infer_body(body);
         self.infer_body(orelse);
     }