[red-knot] Add diagnostic for invalid unpacking

Author: dhruvmanila

crates/red_knot_python_semantic/resources/mdtest/unpacking.md

@@ -61,7 +61,7 @@ reveal_type(c)  # revealed: Literal[4]
 ### Uneven unpacking (1)
 
 ```py
-# TODO: Add diagnostic (there aren't enough values to unpack)
+# error: "Not enough values to unpack (expected 3, got 2)"
 (a, b, c) = (1, 2)
 reveal_type(a)  # revealed: Literal[1]
 reveal_type(b)  # revealed: Literal[2]
@@ -71,7 +71,7 @@ reveal_type(c)  # revealed: Unknown
 ### Uneven unpacking (2)
 
 ```py
-# TODO: Add diagnostic (too many values to unpack)
+# error: "Too many values to unpack (expected 2, got 3)"
 (a, b) = (1, 2, 3)
 reveal_type(a)  # revealed: Literal[1]
 reveal_type(b)  # revealed: Literal[2]
@@ -80,7 +80,7 @@ reveal_type(b)  # revealed: Literal[2]
 ### Starred expression (1)
 
 ```py
-# TODO: Add diagnostic (need more values to unpack)
+# error: "Not enough values to unpack (expected 3 or more, got 2)"
 [a, *b, c, d] = (1, 2)
 reveal_type(a)  # revealed: Literal[1]
 # TODO: Should be list[Any] once support for assigning to starred expression is added
@@ -133,7 +133,7 @@ reveal_type(c)  # revealed: @Todo(starred unpacking)
 ### Starred expression (6)
 
 ```py
-# TODO: Add diagnostic (need more values to unpack)
+# error: "Not enough values to unpack (expected 5 or more, got 1)"
 (a, b, c, *d, e, f) = (1,)
 reveal_type(a)  # revealed: Literal[1]
 reveal_type(b)  # revealed: Unknown
@@ -199,7 +199,7 @@ reveal_type(b)  # revealed: LiteralString
 ### Uneven unpacking (1)
 
 ```py
-# TODO: Add diagnostic (there aren't enough values to unpack)
+# error: "Not enough values to unpack (expected 3, got 2)"
 a, b, c = "ab"
 reveal_type(a)  # revealed: LiteralString
 reveal_type(b)  # revealed: LiteralString
@@ -209,7 +209,7 @@ reveal_type(c)  # revealed: Unknown
 ### Uneven unpacking (2)
 
 ```py
-# TODO: Add diagnostic (too many values to unpack)
+# error: "Too many values to unpack (expected 2, got 3)"
 a, b = "abc"
 reveal_type(a)  # revealed: LiteralString
 reveal_type(b)  # revealed: LiteralString
@@ -218,7 +218,7 @@ reveal_type(b)  # revealed: LiteralString
 ### Starred expression (1)
 
 ```py
-# TODO: Add diagnostic (need more values to unpack)
+# error: "Not enough values to unpack (expected 3 or more, got 2)"
 (a, *b, c, d) = "ab"
 reveal_type(a)  # revealed: LiteralString
 # TODO: Should be list[LiteralString] once support for assigning to starred expression is added
@@ -271,7 +271,7 @@ reveal_type(c)  # revealed: @Todo(starred unpacking)
 ### Unicode
 
 ```py
-# TODO: Add diagnostic (need more values to unpack)
+# error: "Not enough values to unpack (expected 2, got 1)"
 (a, b) = "é"
 
 reveal_type(a)  # revealed: LiteralString
@@ -281,7 +281,7 @@ reveal_type(b)  # revealed: Unknown
 ### Unicode escape (1)
 
 ```py
-# TODO: Add diagnostic (need more values to unpack)
+# error: "Not enough values to unpack (expected 2, got 1)"
 (a, b) = "\u9E6C"
 
 reveal_type(a)  # revealed: LiteralString
@@ -291,7 +291,7 @@ reveal_type(b)  # revealed: Unknown
 ### Unicode escape (2)
 
 ```py
-# TODO: Add diagnostic (need more values to unpack)
+# error: "Not enough values to unpack (expected 2, got 1)"
 (a, b) = "\U0010FFFF"
 
 reveal_type(a)  # revealed: LiteralString
@@ -383,7 +383,8 @@ def _(arg: tuple[int, bytes, int] | tuple[int, int, str, int, bytes]):
 
 ```py
 def _(arg: tuple[int, bytes, int] | tuple[int, int, str, int, bytes]):
-    # TODO: Add diagnostic (too many values to unpack)
+    # error: "Too many values to unpack (expected 2, got 3)"
+    # error: "Too many values to unpack (expected 2, got 5)"
     a, b = arg
     reveal_type(a)  # revealed: int
     reveal_type(b)  # revealed: bytes | int
@@ -393,7 +394,8 @@ def _(arg: tuple[int, bytes, int] | tuple[int, int, str, int, bytes]):
 
 ```py
 def _(arg: tuple[int, bytes] | tuple[int, str]):
-    # TODO: Add diagnostic (there aren't enough values to unpack)
+    # error: "Not enough values to unpack (expected 3, got 2)"
+    # error: "Not enough values to unpack (expected 3, got 2)"
     a, b, c = arg
     reveal_type(a)  # revealed: int
     reveal_type(b)  # revealed: bytes | str
@@ -536,6 +538,7 @@ for a, b in ((1, 2), ("a", "b")):
 # error: "Object of type `Literal[1]` is not iterable"
 # error: "Object of type `Literal[2]` is not iterable"
 # error: "Object of type `Literal[4]` is not iterable"
+# error: "Not enough values to unpack (expected 2, got 1)"
 for a, b in (1, 2, (3, "a"), 4, (5, "b"), "c"):
     reveal_type(a)  # revealed: Unknown | Literal[3, 5] | LiteralString
     reveal_type(b)  # revealed: Unknown | Literal["a", "b"]

crates/red_knot_python_semantic/src/types/display.rs

@@ -324,6 +324,12 @@ impl<'db> TypeArrayDisplay<'db> for Vec<Type<'db>> {
     }
 }
 
+impl<'db> TypeArrayDisplay<'db> for [Type<'db>] {
+    fn display(&self, db: &'db dyn Db) -> DisplayTypeArray {
+        DisplayTypeArray { types: self, db }
+    }
+}
+
 pub(crate) struct DisplayTypeArray<'b, 'db> {
     types: &'b [Type<'db>],
     db: &'db dyn Db,

crates/red_knot_python_semantic/src/types/unpacker.rs

@@ -1,4 +1,5 @@
 use std::borrow::Cow;
+use std::cmp::Ordering;
 
 use rustc_hash::FxHashMap;
 
@@ -11,6 +12,7 @@ use crate::unpack::UnpackValue;
 use crate::Db;
 
 use super::context::{InferContext, WithDiagnostics};
+use super::diagnostic::INVALID_ASSIGNMENT;
 use super::{TupleType, UnionType};
 
 /// Unpacks the value expression type to their respective targets.
@@ -104,9 +106,46 @@ impl<'db> Unpacker<'db> {
                     };
 
                     if let Some(tuple_ty) = ty.into_tuple() {
-                        let tuple_ty_elements = self.tuple_ty_elements(elts, tuple_ty);
+                        let (tuple_ty_elements, has_starred_expression) =
+                            self.tuple_ty_elements(elts, tuple_ty);
 
-                        // TODO: Add diagnostic for length mismatch
+                        match elts.len().cmp(&tuple_ty_elements.len()) {
+                            Ordering::Less => {
+                                self.context.report_lint(
+                                    &INVALID_ASSIGNMENT,
+                                    target.into(),
+                                    format_args!(
+                                        "Too many values to unpack (expected {}, got {})",
+                                        elts.len(),
+                                        tuple_ty_elements.len()
+                                    ),
+                                );
+                            }
+                            Ordering::Greater => {
+                                if has_starred_expression {
+                                    self.context.report_lint(
+                                        &INVALID_ASSIGNMENT,
+                                        target.into(),
+                                        format_args!(
+                                            "Not enough values to unpack (expected {} or more, got {})",
+                                            elts.len() - 1,
+                                            tuple_ty.len(self.db())
+                                        ),
+                                    );
+                                } else {
+                                    self.context.report_lint(
+                                        &INVALID_ASSIGNMENT,
+                                        target.into(),
+                                        format_args!(
+                                            "Not enough values to unpack (expected {}, got {})",
+                                            elts.len(),
+                                            tuple_ty_elements.len()
+                                        ),
+                                    );
+                                }
+                            }
+                            Ordering::Equal => {}
+                        }
 
                         for (index, ty) in tuple_ty_elements.iter().enumerate() {
                             if let Some(element_types) = target_types.get_mut(index) {
@@ -129,6 +168,7 @@ impl<'db> Unpacker<'db> {
                 for (index, element) in elts.iter().enumerate() {
                     // SAFETY: `target_types` is initialized with the same length as `elts`.
                     let element_ty = match target_types[index].as_slice() {
+                        [] if element.is_starred_expr() => todo_type!("starred unpacking"),
                         [] => Type::Unknown,
                         types => UnionType::from_elements(self.db(), types),
                     };
@@ -142,29 +182,39 @@ impl<'db> Unpacker<'db> {
     /// Returns the [`Type`] elements inside the given [`TupleType`] taking into account that there
     /// can be a starred expression in the `elements`.
     fn tuple_ty_elements(
-        &mut self,
+        &self,
         targets: &[ast::Expr],
         tuple_ty: TupleType<'db>,
-    ) -> Cow<'_, [Type<'db>]> {
-        // If there is a starred expression, it will consume all of the entries at that location.
+    ) -> (Cow<'_, [Type<'db>]>, bool) {
+        // If there is a starred expression, it will consume all of the types at that location.
         let Some(starred_index) = targets.iter().position(ast::Expr::is_starred_expr) else {
-            // Otherwise, the types will be unpacked 1-1 to the elements.
-            return Cow::Borrowed(tuple_ty.elements(self.db()).as_ref());
+            // Otherwise, the types will be unpacked 1-1 to the targets.
+            return (Cow::Borrowed(tuple_ty.elements(self.db()).as_ref()), false);
         };
 
         if tuple_ty.len(self.db()) >= targets.len() - 1 {
+            // This branch is only taken when there are enough elements in the tuple type to
+            // combine for the starred expression. So, the arithmetic and indexing operations are
+            // safe to perform.
             let mut element_types = Vec::with_capacity(targets.len());
+
+            // Insert all the elements before the starred expression.
             element_types.extend_from_slice(
                 // SAFETY: Safe because of the length check above.
                 &tuple_ty.elements(self.db())[..starred_index],
             );
 
-            // E.g., in `(a, *b, c, d) = ...`, the index of starred element `b`
-            // is 1 and the remaining elements after that are 2.
+            // The number of target expressions that are remaining after the starred expression.
+            // For example, in `(a, *b, c, d) = ...`, the index of starred element `b` is 1 and the
+            // remaining elements after that are 2.
             let remaining = targets.len() - (starred_index + 1);
-            // This index represents the type of the last element that belongs
-            // to the starred expression, in an exclusive manner.
+
+            // This index represents the position of the last element that belongs to the starred
+            // expression, in an exclusive manner. For example, in `(a, *b, c) = (1, 2, 3, 4)`, the
+            // starred expression `b` will consume the elements `Literal[2]` and `Literal[3]` and
+            // the index value would be 3.
             let starred_end_index = tuple_ty.len(self.db()) - remaining;
+
             // SAFETY: Safe because of the length check above.
             let _starred_element_types =
                 &tuple_ty.elements(self.db())[starred_index..starred_end_index];
@@ -173,19 +223,20 @@ impl<'db> Unpacker<'db> {
             // combine_types(starred_element_types);
             element_types.push(todo_type!("starred unpacking"));
 
+            // Insert the types remaining that aren't consumed by the starred expression.
             element_types.extend_from_slice(
                 // SAFETY: Safe because of the length check above.
                 &tuple_ty.elements(self.db())[starred_end_index..],
             );
-            Cow::Owned(element_types)
-        } else {
+
+            (Cow::Owned(element_types), true)
+        } else if starred_index < tuple_ty.len(self.db()) {
             let mut element_types = tuple_ty.elements(self.db()).to_vec();
-            // Subtract 1 to insert the starred expression type at the correct
-            // index.
-            element_types.resize(targets.len() - 1, Type::Unknown);
             // TODO: This should be `list[Unknown]`
             element_types.insert(starred_index, todo_type!("starred unpacking"));
-            Cow::Owned(element_types)
+            (Cow::Owned(element_types), true)
+        } else {
+            (Cow::Borrowed(tuple_ty.elements(self.db())), true)
         }
     }