mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-30 13:51:37 +00:00
[red-knot] Factor out shared unpacking logic (#16595)
## Summary This PR refactors the common logic for unpacking in assignment, for loops, and with items. ## Test Plan Make sure existing tests pass. --------- Co-authored-by: Dhruv Manilawala <dhruvmanila@gmail.com>
This commit is contained in:
Eric Mark Martin
GitHub
parent
0e48940ea4
commit
78b0b5a3ab
5 changed files with 211 additions and 257 deletions
|
|
@ -38,7 +38,7 @@ use crate::semantic_index::visibility_constraints::{
|
|||
ScopedVisibilityConstraintId, VisibilityConstraintsBuilder,
|
||||
};
|
||||
use crate::semantic_index::SemanticIndex;
|
||||
use crate::unpack::{Unpack, UnpackValue};
|
||||
use crate::unpack::{Unpack, UnpackKind, UnpackPosition, UnpackValue};
|
||||
use crate::Db;
|
||||
|
||||
mod except_handlers;
|
||||
|
|
@ -831,6 +831,64 @@ impl<'db> SemanticIndexBuilder<'db> {
|
|||
debug_assert_eq!(existing_definition, None);
|
||||
}
|
||||
|
||||
/// Add an unpackable assignment for the given [`Unpackable`].
|
||||
///
|
||||
/// This method handles assignments that can contain unpacking like assignment statements,
|
||||
/// for statements, etc.
|
||||
fn add_unpackable_assignment(
|
||||
&mut self,
|
||||
unpackable: &Unpackable<'db>,
|
||||
target: &'db ast::Expr,
|
||||
value: Expression<'db>,
|
||||
) {
|
||||
// We only handle assignments to names and unpackings here, other targets like
|
||||
// attribute and subscript are handled separately as they don't create a new
|
||||
// definition.
|
||||
|
||||
let current_assignment = match target {
|
||||
ast::Expr::List(_) | ast::Expr::Tuple(_) => {
|
||||
let unpack = Some(Unpack::new(
|
||||
self.db,
|
||||
self.file,
|
||||
self.current_scope(),
|
||||
// SAFETY: `target` belongs to the `self.module` tree
|
||||
#[allow(unsafe_code)]
|
||||
unsafe {
|
||||
AstNodeRef::new(self.module.clone(), target)
|
||||
},
|
||||
UnpackValue::new(unpackable.kind(), value),
|
||||
countme::Count::default(),
|
||||
));
|
||||
Some(unpackable.as_current_assignment(unpack))
|
||||
}
|
||||
ast::Expr::Name(_) => Some(unpackable.as_current_assignment(None)),
|
||||
ast::Expr::Attribute(ast::ExprAttribute {
|
||||
value: object,
|
||||
attr,
|
||||
..
|
||||
}) => {
|
||||
self.register_attribute_assignment(
|
||||
object,
|
||||
attr,
|
||||
unpackable.as_attribute_assignment(value),
|
||||
);
|
||||
None
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
if let Some(current_assignment) = current_assignment {
|
||||
self.push_assignment(current_assignment);
|
||||
}
|
||||
|
||||
self.visit_expr(target);
|
||||
|
||||
if current_assignment.is_some() {
|
||||
// Only need to pop in the case where we pushed something
|
||||
self.pop_assignment();
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn build(mut self) -> SemanticIndex<'db> {
|
||||
let module = self.module;
|
||||
self.visit_body(module.suite());
|
||||
|
|
@ -1130,59 +1188,7 @@ where
|
|||
let value = self.add_standalone_expression(&node.value);
|
||||
|
||||
for target in &node.targets {
|
||||
// We only handle assignments to names and unpackings here, other targets like
|
||||
// attribute and subscript are handled separately as they don't create a new
|
||||
// definition.
|
||||
let current_assignment = match target {
|
||||
ast::Expr::List(_) | ast::Expr::Tuple(_) => {
|
||||
Some(CurrentAssignment::Assign {
|
||||
node,
|
||||
first: true,
|
||||
unpack: Some(Unpack::new(
|
||||
self.db,
|
||||
self.file,
|
||||
self.current_scope(),
|
||||
// SAFETY: `target` belongs to the `self.module` tree
|
||||
#[allow(unsafe_code)]
|
||||
unsafe {
|
||||
AstNodeRef::new(self.module.clone(), target)
|
||||
},
|
||||
UnpackValue::Assign(value),
|
||||
countme::Count::default(),
|
||||
)),
|
||||
})
|
||||
}
|
||||
ast::Expr::Name(_) => Some(CurrentAssignment::Assign {
|
||||
node,
|
||||
unpack: None,
|
||||
first: false,
|
||||
}),
|
||||
ast::Expr::Attribute(ast::ExprAttribute {
|
||||
value: object,
|
||||
attr,
|
||||
..
|
||||
}) => {
|
||||
self.register_attribute_assignment(
|
||||
object,
|
||||
attr,
|
||||
AttributeAssignment::Unannotated { value },
|
||||
);
|
||||
|
||||
None
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
if let Some(current_assignment) = current_assignment {
|
||||
self.push_assignment(current_assignment);
|
||||
}
|
||||
|
||||
self.visit_expr(target);
|
||||
|
||||
if current_assignment.is_some() {
|
||||
// Only need to pop in the case where we pushed something
|
||||
self.pop_assignment();
|
||||
}
|
||||
self.add_unpackable_assignment(&Unpackable::Assign(node), target, value);
|
||||
}
|
||||
}
|
||||
ast::Stmt::AnnAssign(node) => {
|
||||
|
|
@ -1373,7 +1379,7 @@ where
|
|||
is_async,
|
||||
..
|
||||
}) => {
|
||||
for item @ ruff_python_ast::WithItem {
|
||||
for item @ ast::WithItem {
|
||||
range: _,
|
||||
context_expr,
|
||||
optional_vars,
|
||||
|
|
@ -1382,55 +1388,14 @@ where
|
|||
self.visit_expr(context_expr);
|
||||
if let Some(optional_vars) = optional_vars.as_deref() {
|
||||
let context_manager = self.add_standalone_expression(context_expr);
|
||||
let current_assignment = match optional_vars {
|
||||
ast::Expr::Tuple(_) | ast::Expr::List(_) => {
|
||||
Some(CurrentAssignment::WithItem {
|
||||
item,
|
||||
first: true,
|
||||
is_async: *is_async,
|
||||
unpack: Some(Unpack::new(
|
||||
self.db,
|
||||
self.file,
|
||||
self.current_scope(),
|
||||
// SAFETY: the node `optional_vars` belongs to the `self.module` tree
|
||||
#[allow(unsafe_code)]
|
||||
unsafe {
|
||||
AstNodeRef::new(self.module.clone(), optional_vars)
|
||||
},
|
||||
UnpackValue::ContextManager(context_manager),
|
||||
countme::Count::default(),
|
||||
)),
|
||||
})
|
||||
}
|
||||
ast::Expr::Name(_) => Some(CurrentAssignment::WithItem {
|
||||
self.add_unpackable_assignment(
|
||||
&Unpackable::WithItem {
|
||||
item,
|
||||
is_async: *is_async,
|
||||
unpack: None,
|
||||
// `false` is arbitrary here---we don't actually use it other than in the actual unpacks
|
||||
first: false,
|
||||
}),
|
||||
ast::Expr::Attribute(ast::ExprAttribute {
|
||||
value: object,
|
||||
attr,
|
||||
..
|
||||
}) => {
|
||||
self.register_attribute_assignment(
|
||||
object,
|
||||
attr,
|
||||
AttributeAssignment::ContextManager { context_manager },
|
||||
);
|
||||
None
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
if let Some(current_assignment) = current_assignment {
|
||||
self.push_assignment(current_assignment);
|
||||
}
|
||||
self.visit_expr(optional_vars);
|
||||
if current_assignment.is_some() {
|
||||
self.pop_assignment();
|
||||
}
|
||||
},
|
||||
optional_vars,
|
||||
context_manager,
|
||||
);
|
||||
}
|
||||
}
|
||||
self.visit_body(body);
|
||||
|
|
@ -1455,52 +1420,7 @@ where
|
|||
|
||||
let pre_loop = self.flow_snapshot();
|
||||
|
||||
let current_assignment = match &**target {
|
||||
ast::Expr::List(_) | ast::Expr::Tuple(_) => Some(CurrentAssignment::For {
|
||||
node: for_stmt,
|
||||
first: true,
|
||||
unpack: Some(Unpack::new(
|
||||
self.db,
|
||||
self.file,
|
||||
self.current_scope(),
|
||||
// SAFETY: the node `target` belongs to the `self.module` tree
|
||||
#[allow(unsafe_code)]
|
||||
unsafe {
|
||||
AstNodeRef::new(self.module.clone(), target)
|
||||
},
|
||||
UnpackValue::Iterable(iter_expr),
|
||||
countme::Count::default(),
|
||||
)),
|
||||
}),
|
||||
ast::Expr::Name(_) => Some(CurrentAssignment::For {
|
||||
node: for_stmt,
|
||||
unpack: None,
|
||||
first: false,
|
||||
}),
|
||||
ast::Expr::Attribute(ast::ExprAttribute {
|
||||
value: object,
|
||||
attr,
|
||||
..
|
||||
}) => {
|
||||
self.register_attribute_assignment(
|
||||
object,
|
||||
attr,
|
||||
AttributeAssignment::Iterable {
|
||||
iterable: iter_expr,
|
||||
},
|
||||
);
|
||||
None
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
|
||||
if let Some(current_assignment) = current_assignment {
|
||||
self.push_assignment(current_assignment);
|
||||
}
|
||||
self.visit_expr(target);
|
||||
if current_assignment.is_some() {
|
||||
self.pop_assignment();
|
||||
}
|
||||
self.add_unpackable_assignment(&Unpackable::For(for_stmt), target, iter_expr);
|
||||
|
||||
let outer_loop = self.push_loop();
|
||||
self.visit_body(body);
|
||||
|
|
@ -1737,18 +1657,13 @@ where
|
|||
|
||||
if is_definition {
|
||||
match self.current_assignment() {
|
||||
Some(CurrentAssignment::Assign {
|
||||
node,
|
||||
first,
|
||||
unpack,
|
||||
}) => {
|
||||
Some(CurrentAssignment::Assign { node, unpack }) => {
|
||||
self.add_definition(
|
||||
symbol,
|
||||
AssignmentDefinitionNodeRef {
|
||||
unpack,
|
||||
value: &node.value,
|
||||
name: name_node,
|
||||
first,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
|
@ -1758,16 +1673,11 @@ where
|
|||
Some(CurrentAssignment::AugAssign(aug_assign)) => {
|
||||
self.add_definition(symbol, aug_assign);
|
||||
}
|
||||
Some(CurrentAssignment::For {
|
||||
node,
|
||||
first,
|
||||
unpack,
|
||||
}) => {
|
||||
Some(CurrentAssignment::For { node, unpack }) => {
|
||||
self.add_definition(
|
||||
symbol,
|
||||
ForStmtDefinitionNodeRef {
|
||||
unpack,
|
||||
first,
|
||||
iterable: &node.iter,
|
||||
name: name_node,
|
||||
is_async: node.is_async,
|
||||
|
|
@ -1793,7 +1703,6 @@ where
|
|||
}
|
||||
Some(CurrentAssignment::WithItem {
|
||||
item,
|
||||
first,
|
||||
is_async,
|
||||
unpack,
|
||||
}) => {
|
||||
|
|
@ -1803,7 +1712,6 @@ where
|
|||
unpack,
|
||||
context_expr: &item.context_expr,
|
||||
name: name_node,
|
||||
first,
|
||||
is_async,
|
||||
},
|
||||
);
|
||||
|
|
@ -1812,13 +1720,11 @@ where
|
|||
}
|
||||
}
|
||||
|
||||
if let Some(
|
||||
CurrentAssignment::Assign { first, .. }
|
||||
| CurrentAssignment::For { first, .. }
|
||||
| CurrentAssignment::WithItem { first, .. },
|
||||
) = self.current_assignment_mut()
|
||||
if let Some(unpack_position) = self
|
||||
.current_assignment_mut()
|
||||
.and_then(CurrentAssignment::unpack_position_mut)
|
||||
{
|
||||
*first = false;
|
||||
*unpack_position = UnpackPosition::Other;
|
||||
}
|
||||
|
||||
walk_expr(self, expr);
|
||||
|
|
@ -1987,20 +1893,10 @@ where
|
|||
ctx: ExprContext::Store,
|
||||
range: _,
|
||||
}) => {
|
||||
if let Some(
|
||||
CurrentAssignment::Assign {
|
||||
unpack: Some(unpack),
|
||||
..
|
||||
}
|
||||
| CurrentAssignment::For {
|
||||
unpack: Some(unpack),
|
||||
..
|
||||
}
|
||||
| CurrentAssignment::WithItem {
|
||||
unpack: Some(unpack),
|
||||
..
|
||||
},
|
||||
) = self.current_assignment()
|
||||
if let Some(unpack) = self
|
||||
.current_assignment()
|
||||
.as_ref()
|
||||
.and_then(CurrentAssignment::unpack)
|
||||
{
|
||||
self.register_attribute_assignment(
|
||||
object,
|
||||
|
|
@ -2075,15 +1971,13 @@ where
|
|||
enum CurrentAssignment<'a> {
|
||||
Assign {
|
||||
node: &'a ast::StmtAssign,
|
||||
first: bool,
|
||||
unpack: Option<Unpack<'a>>,
|
||||
unpack: Option<(UnpackPosition, Unpack<'a>)>,
|
||||
},
|
||||
AnnAssign(&'a ast::StmtAnnAssign),
|
||||
AugAssign(&'a ast::StmtAugAssign),
|
||||
For {
|
||||
node: &'a ast::StmtFor,
|
||||
first: bool,
|
||||
unpack: Option<Unpack<'a>>,
|
||||
unpack: Option<(UnpackPosition, Unpack<'a>)>,
|
||||
},
|
||||
Named(&'a ast::ExprNamed),
|
||||
Comprehension {
|
||||
|
|
@ -2092,12 +1986,37 @@ enum CurrentAssignment<'a> {
|
|||
},
|
||||
WithItem {
|
||||
item: &'a ast::WithItem,
|
||||
first: bool,
|
||||
is_async: bool,
|
||||
unpack: Option<Unpack<'a>>,
|
||||
unpack: Option<(UnpackPosition, Unpack<'a>)>,
|
||||
},
|
||||
}
|
||||
|
||||
impl<'a> CurrentAssignment<'a> {
|
||||
fn unpack(&self) -> Option<Unpack<'a>> {
|
||||
match self {
|
||||
Self::Assign { unpack, .. }
|
||||
| Self::For { unpack, .. }
|
||||
| Self::WithItem { unpack, .. } => unpack.map(|(_, unpack)| unpack),
|
||||
Self::AnnAssign(_)
|
||||
| Self::AugAssign(_)
|
||||
| Self::Named(_)
|
||||
| Self::Comprehension { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn unpack_position_mut(&mut self) -> Option<&mut UnpackPosition> {
|
||||
match self {
|
||||
Self::Assign { unpack, .. }
|
||||
| Self::For { unpack, .. }
|
||||
| Self::WithItem { unpack, .. } => unpack.as_mut().map(|(position, _)| position),
|
||||
Self::AnnAssign(_)
|
||||
| Self::AugAssign(_)
|
||||
| Self::Named(_)
|
||||
| Self::Comprehension { .. } => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> From<&'a ast::StmtAnnAssign> for CurrentAssignment<'a> {
|
||||
fn from(value: &'a ast::StmtAnnAssign) -> Self {
|
||||
Self::AnnAssign(value)
|
||||
|
|
@ -2140,3 +2059,47 @@ impl<'a> CurrentMatchCase<'a> {
|
|||
Self { pattern, index: 0 }
|
||||
}
|
||||
}
|
||||
|
||||
enum Unpackable<'a> {
|
||||
Assign(&'a ast::StmtAssign),
|
||||
For(&'a ast::StmtFor),
|
||||
WithItem {
|
||||
item: &'a ast::WithItem,
|
||||
is_async: bool,
|
||||
},
|
||||
}
|
||||
|
||||
impl<'a> Unpackable<'a> {
|
||||
const fn kind(&self) -> UnpackKind {
|
||||
match self {
|
||||
Unpackable::Assign(_) => UnpackKind::Assign,
|
||||
Unpackable::For(_) => UnpackKind::Iterable,
|
||||
Unpackable::WithItem { .. } => UnpackKind::ContextManager,
|
||||
}
|
||||
}
|
||||
|
||||
fn as_current_assignment(&self, unpack: Option<Unpack<'a>>) -> CurrentAssignment<'a> {
|
||||
let unpack = unpack.map(|unpack| (UnpackPosition::First, unpack));
|
||||
match self {
|
||||
Unpackable::Assign(stmt) => CurrentAssignment::Assign { node: stmt, unpack },
|
||||
Unpackable::For(stmt) => CurrentAssignment::For { node: stmt, unpack },
|
||||
Unpackable::WithItem { item, is_async } => CurrentAssignment::WithItem {
|
||||
item,
|
||||
is_async: *is_async,
|
||||
unpack,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn as_attribute_assignment(&self, expression: Expression<'a>) -> AttributeAssignment<'a> {
|
||||
match self {
|
||||
Unpackable::Assign(_) => AttributeAssignment::Unannotated { value: expression },
|
||||
Unpackable::For(_) => AttributeAssignment::Iterable {
|
||||
iterable: expression,
|
||||
},
|
||||
Unpackable::WithItem { .. } => AttributeAssignment::ContextManager {
|
||||
context_manager: expression,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue