[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>

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -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
+                    self.add_unpackable_assignment(&Unpackable::Assign(node), target, value);
-                    // 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();
-                    }
                 }
             }
             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 {
+                        self.add_unpackable_assignment(
-                            ast::Expr::Tuple(_) | ast::Expr::List(_) => {
+                            &Unpackable::WithItem {
-                                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 {
                                 item,
                                 is_async: *is_async,
-                                unpack: None,
+                            },
-                                // `false` is arbitrary here---we don't actually use it other than in the actual unpacks
+                            optional_vars,
-                                first: false,
+                            context_manager,
-                            }),
+                        );
-                            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();
-                        }
                     }
                 }
                 self.visit_body(body);
@@ -1455,52 +1420,7 @@ where
 
                 let pre_loop = self.flow_snapshot();
 
-                let current_assignment = match &**target {
+                self.add_unpackable_assignment(&Unpackable::For(for_stmt), target, iter_expr);
-                    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();
-                }
 
                 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 {
+                        Some(CurrentAssignment::Assign { node, unpack }) => {
-                            node,
-                            first,
-                            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 {
+                        Some(CurrentAssignment::For { node, unpack }) => {
-                            node,
-                            first,
-                            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(
+                if let Some(unpack_position) = self
-                    CurrentAssignment::Assign { first, .. }
+                    .current_assignment_mut()
-                    | CurrentAssignment::For { first, .. }
+                    .and_then(CurrentAssignment::unpack_position_mut)
-                    | CurrentAssignment::WithItem { first, .. },
-                ) = self.current_assignment_mut()
                 {
-                    *first = false;
+                    *unpack_position = UnpackPosition::Other;
                 }
 
                 walk_expr(self, expr);
@@ -1987,20 +1893,10 @@ where
                 ctx: ExprContext::Store,
                 range: _,
             }) => {
-                if let Some(
+                if let Some(unpack) = self
-                    CurrentAssignment::Assign {
+                    .current_assignment()
-                        unpack: Some(unpack),
+                    .as_ref()
-                        ..
+                    .and_then(CurrentAssignment::unpack)
-                    }
-                    | CurrentAssignment::For {
-                        unpack: Some(unpack),
-                        ..
-                    }
-                    | CurrentAssignment::WithItem {
-                        unpack: Some(unpack),
-                        ..
-                    },
-                ) = self.current_assignment()
                 {
                     self.register_attribute_assignment(
                         object,
@@ -2075,15 +1971,13 @@ where
 enum CurrentAssignment<'a> {
     Assign {
         node: &'a ast::StmtAssign,
-        first: bool,
+        unpack: Option<(UnpackPosition, Unpack<'a>)>,
-        unpack: Option<Unpack<'a>>,
     },
     AnnAssign(&'a ast::StmtAnnAssign),
     AugAssign(&'a ast::StmtAugAssign),
     For {
         node: &'a ast::StmtFor,
-        first: bool,
+        unpack: Option<(UnpackPosition, Unpack<'a>)>,
-        unpack: Option<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,
+            },
+        }
+    }
+}

crates/red_knot_python_semantic/src/semantic_index/definition.rs

@@ -8,7 +8,7 @@ use ruff_text_size::{Ranged, TextRange};
 use crate::ast_node_ref::AstNodeRef;
 use crate::node_key::NodeKey;
 use crate::semantic_index::symbol::{FileScopeId, ScopeId, ScopedSymbolId};
-use crate::unpack::Unpack;
+use crate::unpack::{Unpack, UnpackPosition};
 use crate::Db;
 
 /// A definition of a symbol.
@@ -239,27 +239,24 @@ pub(crate) struct ImportFromDefinitionNodeRef<'a> {
 
 #[derive(Copy, Clone, Debug)]
 pub(crate) struct AssignmentDefinitionNodeRef<'a> {
-    pub(crate) unpack: Option<Unpack<'a>>,
+    pub(crate) unpack: Option<(UnpackPosition, Unpack<'a>)>,
     pub(crate) value: &'a ast::Expr,
     pub(crate) name: &'a ast::ExprName,
-    pub(crate) first: bool,
 }
 
 #[derive(Copy, Clone, Debug)]
 pub(crate) struct WithItemDefinitionNodeRef<'a> {
-    pub(crate) unpack: Option<Unpack<'a>>,
+    pub(crate) unpack: Option<(UnpackPosition, Unpack<'a>)>,
     pub(crate) context_expr: &'a ast::Expr,
     pub(crate) name: &'a ast::ExprName,
-    pub(crate) first: bool,
     pub(crate) is_async: bool,
 }
 
 #[derive(Copy, Clone, Debug)]
 pub(crate) struct ForStmtDefinitionNodeRef<'a> {
-    pub(crate) unpack: Option<Unpack<'a>>,
+    pub(crate) unpack: Option<(UnpackPosition, Unpack<'a>)>,
     pub(crate) iterable: &'a ast::Expr,
     pub(crate) name: &'a ast::ExprName,
-    pub(crate) first: bool,
     pub(crate) is_async: bool,
 }
 
@@ -332,12 +329,10 @@ impl<'db> DefinitionNodeRef<'db> {
                 unpack,
                 value,
                 name,
-                first,
             }) => DefinitionKind::Assignment(AssignmentDefinitionKind {
                 target: TargetKind::from(unpack),
                 value: AstNodeRef::new(parsed.clone(), value),
                 name: AstNodeRef::new(parsed, name),
-                first,
             }),
             DefinitionNodeRef::AnnotatedAssignment(assign) => {
                 DefinitionKind::AnnotatedAssignment(AstNodeRef::new(parsed, assign))
@@ -349,13 +344,11 @@ impl<'db> DefinitionNodeRef<'db> {
                 unpack,
                 iterable,
                 name,
-                first,
                 is_async,
             }) => DefinitionKind::For(ForStmtDefinitionKind {
                 target: TargetKind::from(unpack),
                 iterable: AstNodeRef::new(parsed.clone(), iterable),
                 name: AstNodeRef::new(parsed, name),
-                first,
                 is_async,
             }),
             DefinitionNodeRef::Comprehension(ComprehensionDefinitionNodeRef {
@@ -382,13 +375,11 @@ impl<'db> DefinitionNodeRef<'db> {
                 unpack,
                 context_expr,
                 name,
-                first,
                 is_async,
             }) => DefinitionKind::WithItem(WithItemDefinitionKind {
                 target: TargetKind::from(unpack),
                 context_expr: AstNodeRef::new(parsed.clone(), context_expr),
                 name: AstNodeRef::new(parsed, name),
-                first,
                 is_async,
             }),
             DefinitionNodeRef::MatchPattern(MatchPatternDefinitionNodeRef {
@@ -451,7 +442,6 @@ impl<'db> DefinitionNodeRef<'db> {
                 value: _,
                 unpack: _,
                 name,
-                first: _,
             }) => name.into(),
             Self::AnnotatedAssignment(node) => node.into(),
             Self::AugmentedAssignment(node) => node.into(),
@@ -459,7 +449,6 @@ impl<'db> DefinitionNodeRef<'db> {
                 unpack: _,
                 iterable: _,
                 name,
-                first: _,
                 is_async: _,
             }) => name.into(),
             Self::Comprehension(ComprehensionDefinitionNodeRef { target, .. }) => target.into(),
@@ -469,7 +458,6 @@ impl<'db> DefinitionNodeRef<'db> {
             Self::WithItem(WithItemDefinitionNodeRef {
                 unpack: _,
                 context_expr: _,
-                first: _,
                 is_async: _,
                 name,
             }) => name.into(),
@@ -652,14 +640,14 @@ impl DefinitionKind<'_> {
 
 #[derive(Copy, Clone, Debug, PartialEq, Hash)]
 pub(crate) enum TargetKind<'db> {
-    Sequence(Unpack<'db>),
+    Sequence(UnpackPosition, Unpack<'db>),
     Name,
 }
 
-impl<'db> From<Option<Unpack<'db>>> for TargetKind<'db> {
+impl<'db> From<Option<(UnpackPosition, Unpack<'db>)>> for TargetKind<'db> {
-    fn from(value: Option<Unpack<'db>>) -> Self {
+    fn from(value: Option<(UnpackPosition, Unpack<'db>)>) -> Self {
         match value {
-            Some(unpack) => TargetKind::Sequence(unpack),
+            Some((unpack_position, unpack)) => TargetKind::Sequence(unpack_position, unpack),
             None => TargetKind::Name,
         }
     }
@@ -780,7 +768,6 @@ pub struct AssignmentDefinitionKind<'db> {
     target: TargetKind<'db>,
     value: AstNodeRef<ast::Expr>,
     name: AstNodeRef<ast::ExprName>,
-    first: bool,
 }
 
 impl<'db> AssignmentDefinitionKind<'db> {
@@ -795,10 +782,6 @@ impl<'db> AssignmentDefinitionKind<'db> {
     pub(crate) fn name(&self) -> &ast::ExprName {
         self.name.node()
     }
-
-    pub(crate) fn is_first(&self) -> bool {
-        self.first
-    }
 }
 
 #[derive(Clone, Debug)]
@@ -806,7 +789,6 @@ pub struct WithItemDefinitionKind<'db> {
     target: TargetKind<'db>,
     context_expr: AstNodeRef<ast::Expr>,
     name: AstNodeRef<ast::ExprName>,
-    first: bool,
     is_async: bool,
 }
 
@@ -823,10 +805,6 @@ impl<'db> WithItemDefinitionKind<'db> {
         self.name.node()
     }
 
-    pub(crate) const fn is_first(&self) -> bool {
-        self.first
-    }
-
     pub(crate) const fn is_async(&self) -> bool {
         self.is_async
     }
@@ -837,7 +815,6 @@ pub struct ForStmtDefinitionKind<'db> {
     target: TargetKind<'db>,
     iterable: AstNodeRef<ast::Expr>,
     name: AstNodeRef<ast::ExprName>,
-    first: bool,
     is_async: bool,
 }
 
@@ -854,10 +831,6 @@ impl<'db> ForStmtDefinitionKind<'db> {
         self.name.node()
     }
 
-    pub(crate) const fn is_first(&self) -> bool {
-        self.first
-    }
-
     pub(crate) const fn is_async(&self) -> bool {
         self.is_async
     }

crates/red_knot_python_semantic/src/types/infer.rs

@@ -84,7 +84,7 @@ use crate::types::{
     UnionType,
 };
 use crate::types::{CallableType, GeneralCallableType, Signature};
-use crate::unpack::Unpack;
+use crate::unpack::{Unpack, UnpackPosition};
 use crate::util::subscript::{PyIndex, PySlice};
 use crate::Db;
 
@@ -1823,10 +1823,10 @@ impl<'db> TypeInferenceBuilder<'db> {
             todo_type!("async `with` statement")
         } else {
             match with_item.target() {
-                TargetKind::Sequence(unpack) => {
+                TargetKind::Sequence(unpack_position, unpack) => {
                     let unpacked = infer_unpack_types(self.db(), unpack);
                     let name_ast_id = name.scoped_expression_id(self.db(), self.scope());
-                    if with_item.is_first() {
+                    if unpack_position == UnpackPosition::First {
                         self.context.extend(unpacked);
                     }
                     unpacked.expression_type(name_ast_id)
@@ -2627,11 +2627,11 @@ impl<'db> TypeInferenceBuilder<'db> {
         let value_ty = self.infer_standalone_expression(value);
 
         let mut target_ty = match assignment.target() {
-            TargetKind::Sequence(unpack) => {
+            TargetKind::Sequence(unpack_position, unpack) => {
                 let unpacked = infer_unpack_types(self.db(), unpack);
                 // Only copy the diagnostics if this is the first assignment to avoid duplicating the
                 // unpack assignments.
-                if assignment.is_first() {
+                if unpack_position == UnpackPosition::First {
                     self.context.extend(unpacked);
                 }
 
@@ -2929,9 +2929,9 @@ impl<'db> TypeInferenceBuilder<'db> {
             todo_type!("async iterables/iterators")
         } else {
             match for_stmt.target() {
-                TargetKind::Sequence(unpack) => {
+                TargetKind::Sequence(unpack_position, unpack) => {
                     let unpacked = infer_unpack_types(self.db(), unpack);
-                    if for_stmt.is_first() {
+                    if unpack_position == UnpackPosition::First {
                         self.context.extend(unpacked);
                     }
                     let name_ast_id = name.scoped_expression_id(self.db(), self.scope());

crates/red_knot_python_semantic/src/types/unpacker.rs

@@ -8,7 +8,7 @@ use ruff_python_ast::{self as ast, AnyNodeRef};
 use crate::semantic_index::ast_ids::{HasScopedExpressionId, ScopedExpressionId};
 use crate::semantic_index::symbol::ScopeId;
 use crate::types::{infer_expression_types, todo_type, Type, TypeCheckDiagnostics};
-use crate::unpack::UnpackValue;
+use crate::unpack::{UnpackKind, UnpackValue};
 use crate::Db;
 
 use super::context::{InferContext, WithDiagnostics};
@@ -45,30 +45,27 @@ impl<'db> Unpacker<'db> {
         let value_type = infer_expression_types(self.db(), value.expression())
             .expression_type(value.scoped_expression_id(self.db(), self.scope));
 
-        let value_type = match value {
+        let value_type = match value.kind() {
-            UnpackValue::Assign(expression) => {
+            UnpackKind::Assign => {
                 if self.context.in_stub()
-                    && expression.node_ref(self.db()).is_ellipsis_literal_expr()
+                    && value
+                        .expression()
+                        .node_ref(self.db())
+                        .is_ellipsis_literal_expr()
                 {
                     Type::unknown()
                 } else {
                     value_type
                 }
             }
-            UnpackValue::Iterable(_) => value_type.try_iterate(self.db()).unwrap_or_else(|err| {
+            UnpackKind::Iterable => value_type.try_iterate(self.db()).unwrap_or_else(|err| {
                 err.report_diagnostic(&self.context, value_type, value.as_any_node_ref(self.db()));
                 err.fallback_element_type(self.db())
             }),
-            UnpackValue::ContextManager(_) => {
+            UnpackKind::ContextManager => value_type.try_enter(self.db()).unwrap_or_else(|err| {
-                value_type.try_enter(self.db()).unwrap_or_else(|err| {
+                err.report_diagnostic(&self.context, value_type, value.as_any_node_ref(self.db()));
-                    err.report_diagnostic(
+                err.fallback_enter_type(self.db())
-                        &self.context,
+            }),
-                        value_type,
-                        value.as_any_node_ref(self.db()),
-                    );
-                    err.fallback_enter_type(self.db())
-                })
-            }
         };
 
         self.unpack_inner(target, value.as_any_node_ref(self.db()), value_type);

crates/red_knot_python_semantic/src/unpack.rs

@@ -60,23 +60,21 @@ impl<'db> Unpack<'db> {
 
 /// The expression that is being unpacked.
 #[derive(Clone, Copy, Debug, Hash, salsa::Update)]
-pub(crate) enum UnpackValue<'db> {
+pub(crate) struct UnpackValue<'db> {
-    /// An iterable expression like the one in a `for` loop or a comprehension.
+    /// The kind of unpack expression
-    Iterable(Expression<'db>),
+    kind: UnpackKind,
-    /// An context manager expression like the one in a `with` statement.
+    /// The expression we are unpacking
-    ContextManager(Expression<'db>),
+    expression: Expression<'db>,
-    /// An expression that is being assigned to a target.
-    Assign(Expression<'db>),
 }
 
 impl<'db> UnpackValue<'db> {
+    pub(crate) fn new(kind: UnpackKind, expression: Expression<'db>) -> Self {
+        Self { kind, expression }
+    }
+
     /// Returns the underlying [`Expression`] that is being unpacked.
     pub(crate) const fn expression(self) -> Expression<'db> {
-        match self {
+        self.expression
-            UnpackValue::Assign(expr)
-            | UnpackValue::Iterable(expr)
-            | UnpackValue::ContextManager(expr) => expr,
-        }
     }
 
     /// Returns the [`ScopedExpressionId`] of the underlying expression.
@@ -94,4 +92,27 @@ impl<'db> UnpackValue<'db> {
     pub(crate) fn as_any_node_ref(self, db: &'db dyn Db) -> AnyNodeRef<'db> {
         self.expression().node_ref(db).node().into()
     }
+
+    pub(crate) const fn kind(self) -> UnpackKind {
+        self.kind
+    }
+}
+
+#[derive(Clone, Copy, Debug, Hash, salsa::Update)]
+pub(crate) enum UnpackKind {
+    /// An iterable expression like the one in a `for` loop or a comprehension.
+    Iterable,
+    /// An context manager expression like the one in a `with` statement.
+    ContextManager,
+    /// An expression that is being assigned to a target.
+    Assign,
+}
+
+/// The position of the target element in an unpacking.
+#[derive(Clone, Copy, Debug, Hash, PartialEq, salsa::Update)]
+pub(crate) enum UnpackPosition {
+    /// The target element is in the first position of the unpacking.
+    First,
+    /// The target element is in the position other than the first position of the unpacking.
+    Other,
 }