language-servers/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-09-29 13:25:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

[red-knot] Infer target types for unpacked tuple assignment (#13316)

Browse source 
## Summary

This PR adds support for unpacking tuple expression in an assignment
statement where the target expression can be a tuple or a list (the
allowed sequence targets).

The implementation introduces a new `infer_assignment_target` which can
then be used for other targets like the ones in for loops as well. This
delegates it to the `infer_definition`. The final implementation uses a
recursive function that visits the target expression in source order and
compares the variable node that corresponds to the definition. At the
same time, it keeps track of where it is on the assignment value type.

The logic also accounts for the number of elements on both sides such
that it matches even if there's a gap in between. For example, if
there's a starred expression like `(a, *b, c) = (1, 2, 3)`, then the
type of `a` will be `Literal[1]` and the type of `b` will be
`Literal[2]`.

There are a couple of follow-ups that can be done:
* Use this logic for other target positions like `for` loop
* Add diagnostics for mis-match length between LHS and RHS

## Test Plan

Add various test cases using the new markdown test framework.
Validate that existing test cases pass.

---------

Co-authored-by: Carl Meyer <carl@astral.sh>
This commit is contained in:
Dhruv Manilawala 2024-10-16 00:37:11 +05:30 • committed by GitHub
parent d77480768d
commit b16f665a81
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 525 additions and 46 deletions
Download patch file Download diff file Expand all files Collapse all files

43
crates/red_knot_python_semantic/src/semantic_index/builder.rs
View file

@ -28,8 +28,8 @@ use crate::Db;
use super::constraint::{Constraint, PatternConstraint};
use super::definition::{
DefinitionCategory, ExceptHandlerDefinitionNodeRef, MatchPatternDefinitionNodeRef,
WithItemDefinitionNodeRef,
AssignmentKind, DefinitionCategory, ExceptHandlerDefinitionNodeRef,
MatchPatternDefinitionNodeRef, WithItemDefinitionNodeRef,
};
pub(super) struct SemanticIndexBuilder<'db> {
@ -566,11 +566,22 @@ where
debug_assert!(self.current_assignment.is_none());
self.visit_expr(&node.value);
self.add_standalone_expression(&node.value);
self.current_assignment = Some(node.into());
for target in &node.targets {
for (target_index, target) in node.targets.iter().enumerate() {
let kind = match target {
ast::Expr::List(_) | ast::Expr::Tuple(_) => Some(AssignmentKind::Sequence),
ast::Expr::Name(_) => Some(AssignmentKind::Name),
_ => None,
};
if let Some(kind) = kind {
self.current_assignment = Some(CurrentAssignment::Assign {
assignment: node,
target_index,
kind,
});
}
self.visit_expr(target);
self.current_assignment = None;
}
self.current_assignment = None;
}
ast::Stmt::AnnAssign(node) => {
debug_assert!(self.current_assignment.is_none());
@ -815,12 +826,18 @@ where
let symbol = self.add_symbol(id.clone());
if is_definition {
match self.current_assignment {
Some(CurrentAssignment::Assign(assignment)) => {
Some(CurrentAssignment::Assign {
assignment,
target_index,
kind,
}) => {
self.add_definition(
symbol,
AssignmentDefinitionNodeRef {
assignment,
target: name_node,
target_index,
name: name_node,
kind,
},
);
}
@ -1045,7 +1062,11 @@ where
#[derive(Copy, Clone, Debug)]
enum CurrentAssignment<'a> {
Assign(&'a ast::StmtAssign),
Assign {
assignment: &'a ast::StmtAssign,
target_index: usize,
kind: AssignmentKind,
},
AnnAssign(&'a ast::StmtAnnAssign),
AugAssign(&'a ast::StmtAugAssign),
For(&'a ast::StmtFor),
@ -1057,12 +1078,6 @@ enum CurrentAssignment<'a> {
WithItem(&'a ast::WithItem),
}
impl<'a> From<&'a ast::StmtAssign> for CurrentAssignment<'a> {
fn from(value: &'a ast::StmtAssign) -> Self {
Self::Assign(value)
}
}
impl<'a> From<&'a ast::StmtAnnAssign> for CurrentAssignment<'a> {
fn from(value: &'a ast::StmtAnnAssign) -> Self {
Self::AnnAssign(value)