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 07c51c8
Show file tree
Hide file tree
Showing 3 changed files with 77 additions and 21 deletions.
28 changes: 23 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,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
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
62 changes: 46 additions & 16 deletions crates/red_knot_python_semantic/src/types/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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);
}
}
Expand Down Expand Up @@ -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);
}
Expand Down

0 comments on commit 07c51c8

Please sign in to comment.