Skip to content

Commit

Permalink
Also check assignment target in unpack assignments and for loops
Browse files Browse the repository at this point in the history
  • Loading branch information
sharkdp committed Jan 20, 2025
1 parent 1ee40b7 commit 51cde27
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 17 deletions.
25 changes: 20 additions & 5 deletions crates/red_knot_python_semantic/resources/mdtest/attributes.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 8 additions & 0 deletions crates/red_knot_python_semantic/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
54 changes: 42 additions & 12 deletions crates/red_knot_python_semantic/src/types/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<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(
Expand All @@ -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);
}
}
Expand Down Expand Up @@ -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);
}
Expand Down

0 comments on commit 51cde27

Please sign in to comment.