From 51cde274dd3c0b25513fea5f075522f0e581a833 Mon Sep 17 00:00:00 2001 From: David Peter Date: Mon, 20 Jan 2025 13:09:17 +0100 Subject: [PATCH] Also check assignment target in unpack assignments and for loops --- .../resources/mdtest/attributes.md | 25 +++++++-- crates/red_knot_python_semantic/src/types.rs | 8 +++ .../src/types/infer.rs | 54 ++++++++++++++----- 3 files changed, 70 insertions(+), 17 deletions(-) diff --git a/crates/red_knot_python_semantic/resources/mdtest/attributes.md b/crates/red_knot_python_semantic/resources/mdtest/attributes.md index 2178c17a00a55c..82230cfa329836 100644 --- a/crates/red_knot_python_semantic/resources/mdtest/attributes.md +++ b/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 diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs index 49747b105645ef..2a21f4819888f6 100644 --- a/crates/red_knot_python_semantic/src/types.rs +++ b/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)] diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs index 99125e4fe4905f..b5b0b1f91f457b 100644 --- a/crates/red_knot_python_semantic/src/types/infer.rs +++ b/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 in [`infer_definition`]. /// /// # 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>, + ) { 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> = + 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); }