Also check assignment target in unpack assignments and for loops

Author: sharkdp

crates/red_knot_python_semantic/resources/mdtest/attributes.md

@@ -439,17 +439,35 @@ 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)
+
+# TODO: this should be an error, but we do not understand lists yet.
+[_, mod.global_symbol] = [1, 2]
+
+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

@@ -1969,28 +1969,51 @@ impl<'db> TypeInferenceBuilder<'db> {
         } = assignment;
 
         for target in targets {
-            self.infer_target(target, value);
+            self.infer_target(target, value, |_, ty| ty);
         }
     }
 
-    /// 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.
+    ///
+    /// The `to_assigned_ty` function is used to convert the inferred type of the `value` expression
+    /// to the type that is eventually assigned to the `target`.
     ///
     /// # Panics
     ///
     /// If the `value` is not a standalone expression.
-    fn infer_target(&mut self, target: &ast::Expr, value: &ast::Expr) {
+    fn infer_target<F>(&mut self, target: &ast::Expr, value: &ast::Expr, to_assigned_ty: F)
+    where
+        F: Fn(&'db dyn Db, Type<'db>) -> Type<'db>,
+    {
+        let assigned_ty = match target {
+            ast::Expr::Name(_) => None,
+            _ => {
+                let value_ty = self.infer_standalone_expression(value);
+
+                Some(to_assigned_ty(self.db(), value_ty))
+            }
+        };
+        self.infer_target_impl(target, assigned_ty);
+    }
+
+    fn infer_target_impl(&mut self, target: &ast::Expr, assigned_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 assigned_tys = match assigned_ty {
+                    Some(Type::Tuple(tuple)) => {
+                        Either::Left(tuple.elements(self.db()).into_iter().copied())
+                    }
+                    Some(_) | None => Either::Right(std::iter::empty()),
+                };
+
                 for element in elts {
-                    self.infer_target(element, value);
-                }
-                if elts.is_empty() {
-                    self.infer_standalone_expression(value);
+                    self.infer_target_impl(element, assigned_tys.next());
                 }
             }
             ast::Expr::Attribute(
@@ -1999,18 +2022,22 @@ impl<'db> TypeInferenceBuilder<'db> {
                     ..
                 },
             ) => {
-                let lhs_ty = self.infer_attribute_expression(lhs_expr);
-                self.store_expression_type(target, lhs_ty);
+                let attribute_expr_ty = self.infer_attribute_expression(lhs_expr);
+                self.store_expression_type(target, attribute_expr_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(assigned_ty) = assigned_ty {
+                    if !assigned_ty.is_assignable_to(self.db(), attribute_expr_ty) {
+                        report_invalid_assignment(
+                            &self.context,
+                            target.into(),
+                            attribute_expr_ty,
+                            assigned_ty,
+                        );
+                    }
                 }
             }
             _ => {
                 // TODO: Remove this once we handle all possible assignment targets.
-                self.infer_standalone_expression(value);
                 self.infer_expression(target);
             }
         }
@@ -2279,7 +2306,10 @@ impl<'db> TypeInferenceBuilder<'db> {
             is_async: _,
         } = for_statement;
 
-        self.infer_target(target, iter);
+        self.infer_target(target, iter, |db, iter_ty| {
+            iter_ty.iterate(db).unwrap_without_diagnostic()
+        });
+
         self.infer_body(body);
         self.infer_body(orelse);
     }