[red-knot] Infer tuple types from annotations (#13943)

## Summary

This PR adds support for heterogenous `tuple` annotations to red-knot.

It does the following:
- Extends `infer_type_expression` so that it understands tuple
annotations
- Changes `infer_type_expression` so that `ExprStarred` nodes in type
annotations are inferred as `Todo` rather than `Unknown` (they're valid
in PEP-646 tuple annotations)
- Extends `Type::is_subtype_of` to understand when one heterogenous
tuple type can be understood to be a subtype of another (without this
change, the PR would have introduced new false-positive errors to some
existing mdtests).

crates/red_knot_python_semantic/resources/mdtest/assignment/annotations.md

@@ -23,6 +23,56 @@ x: int
 x = "foo"  # error: [invalid-assignment] "Object of type `Literal["foo"]` is not assignable to `int`"
 ```
 
+## Tuple annotations are understood
+
+```py path=module.py
+from typing_extensions import Unpack
+
+a: tuple[()] = ()
+b: tuple[int] = (42,)
+c: tuple[str, int] = ("42", 42)
+d: tuple[tuple[str, str], tuple[int, int]] = (("foo", "foo"), (42, 42))
+e: tuple[str, ...] = ()
+f: tuple[str, *tuple[int, ...], bytes] = ("42", b"42")
+g: tuple[str, Unpack[tuple[int, ...]], bytes] = ("42", b"42")
+h: tuple[list[int], list[int]] = ([], [])
+i: tuple[str | int, str | int] = (42, 42)
+j: tuple[str | int] = (42,)
+```
+
+```py path=script.py
+from module import a, b, c, d, e, f, g, h, i, j
+
+reveal_type(a)  # revealed: tuple[()]
+reveal_type(b)  # revealed: tuple[int]
+reveal_type(c)  # revealed: tuple[str, int]
+reveal_type(d)  # revealed: tuple[tuple[str, str], tuple[int, int]]
+
+# TODO: homogenous tuples, PEP-646 tuples
+reveal_type(e)  # revealed: @Todo
+reveal_type(f)  # revealed: @Todo
+reveal_type(g)  # revealed: @Todo
+
+# TODO: support more kinds of type expressions in annotations
+reveal_type(h)  # revealed: @Todo
+
+reveal_type(i)  # revealed: tuple[str | int, str | int]
+reveal_type(j)  # revealed: tuple[str | int]
+```
+
+## Incorrect tuple assignments are complained about
+
+```py
+# error: [invalid-assignment] "Object of type `tuple[Literal[1], Literal[2]]` is not assignable to `tuple[()]`"
+a: tuple[()] = (1, 2)
+
+# error: [invalid-assignment] "Object of type `tuple[Literal["foo"]]` is not assignable to `tuple[int]`"
+b: tuple[int] = ("foo",)
+
+# error: [invalid-assignment] "Object of type `tuple[list, Literal["foo"]]` is not assignable to `tuple[str | int, str]`"
+c: tuple[str | int, str] = ([], "foo")
+```
+
 ## PEP-604 annotations are supported
 
 ```py

crates/red_knot_python_semantic/resources/mdtest/generics.md

@@ -16,10 +16,11 @@ class MyBox[T]:
     def __init__(self, data: T):
         self.data = data
 
-# TODO not error (should be subscriptable)
+box: MyBox[int] = MyBox(5)
-box: MyBox[int] = MyBox(5)  # error: [non-subscriptable]
+
-# TODO error differently (str and int don't unify)
+# TODO should emit a diagnostic here (str is not assignable to int)
-wrong_innards: MyBox[int] = MyBox("five")  # error: [non-subscriptable]
+wrong_innards: MyBox[int] = MyBox("five")
+
 # TODO reveal int
 reveal_type(box.data)  # revealed: @Todo
 

crates/red_knot_python_semantic/src/types.rs

@@ -469,6 +469,16 @@ impl<'db> Type<'db> {
             {
                 true
             }
+            (Type::Tuple(self_tuple), Type::Tuple(target_tuple)) => {
+                let self_elements = self_tuple.elements(db);
+                let target_elements = target_tuple.elements(db);
+                self_elements.len() == target_elements.len()
+                    && self_elements.iter().zip(target_elements).all(
+                        |(self_element, target_element)| {
+                            self_element.is_subtype_of(db, *target_element)
+                        },
+                    )
+            }
             (Type::ClassLiteral(..), Type::Instance(class))
                 if class.is_known(db, KnownClass::Type) =>
             {
@@ -504,6 +514,16 @@ impl<'db> Type<'db> {
                 .elements(db)
                 .iter()
                 .any(|&elem_ty| ty.is_assignable_to(db, elem_ty)),
+            (Type::Tuple(self_tuple), Type::Tuple(target_tuple)) => {
+                let self_elements = self_tuple.elements(db);
+                let target_elements = target_tuple.elements(db);
+                self_elements.len() == target_elements.len()
+                    && self_elements.iter().zip(target_elements).all(
+                        |(self_element, target_element)| {
+                            self_element.is_assignable_to(db, *target_element)
+                        },
+                    )
+            }
             // TODO other types containing gradual forms (e.g. generics containing Any/Unknown)
             _ => self.is_subtype_of(db, target),
         }
@@ -1887,6 +1907,7 @@ mod tests {
         Unknown,
         None,
         Any,
+        Todo,
         IntLiteral(i64),
         BooleanLiteral(bool),
         StringLiteral(&'static str),
@@ -1905,6 +1926,7 @@ mod tests {
                 Ty::Unknown => Type::Unknown,
                 Ty::None => Type::None,
                 Ty::Any => Type::Any,
+                Ty::Todo => Type::Todo,
                 Ty::IntLiteral(n) => Type::IntLiteral(n),
                 Ty::StringLiteral(s) => Type::StringLiteral(StringLiteralType::new(db, s)),
                 Ty::BooleanLiteral(b) => Type::BooleanLiteral(b),
@@ -1947,6 +1969,8 @@ mod tests {
     #[test_case(Ty::IntLiteral(1), Ty::Union(vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str")]))]
     #[test_case(Ty::IntLiteral(1), Ty::Union(vec![Ty::Unknown, Ty::BuiltinInstance("str")]))]
     #[test_case(Ty::Union(vec![Ty::IntLiteral(1), Ty::IntLiteral(2)]), Ty::Union(vec![Ty::IntLiteral(1), Ty::IntLiteral(2)]))]
+    #[test_case(Ty::Tuple(vec![Ty::Todo]), Ty::Tuple(vec![Ty::IntLiteral(2)]))]
+    #[test_case(Ty::Tuple(vec![Ty::IntLiteral(2)]), Ty::Tuple(vec![Ty::Todo]))]
     fn is_assignable_to(from: Ty, to: Ty) {
         let db = setup_db();
         assert!(from.into_type(&db).is_assignable_to(&db, to.into_type(&db)));
@@ -1981,6 +2005,11 @@ mod tests {
     #[test_case(Ty::Union(vec![Ty::BuiltinInstance("str"), Ty::BuiltinInstance("int")]), Ty::BuiltinInstance("object"))]
     #[test_case(Ty::Union(vec![Ty::IntLiteral(1), Ty::IntLiteral(2)]), Ty::Union(vec![Ty::IntLiteral(1), Ty::IntLiteral(2), Ty::IntLiteral(3)]))]
     #[test_case(Ty::BuiltinInstance("TypeError"), Ty::BuiltinInstance("Exception"))]
+    #[test_case(Ty::Tuple(vec![]), Ty::Tuple(vec![]))]
+    #[test_case(Ty::Tuple(vec![Ty::IntLiteral(42)]), Ty::Tuple(vec![Ty::BuiltinInstance("int")]))]
+    #[test_case(Ty::Tuple(vec![Ty::IntLiteral(42), Ty::StringLiteral("foo")]), Ty::Tuple(vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str")]))]
+    #[test_case(Ty::Tuple(vec![Ty::BuiltinInstance("int"), Ty::StringLiteral("foo")]), Ty::Tuple(vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str")]))]
+    #[test_case(Ty::Tuple(vec![Ty::IntLiteral(42), Ty::BuiltinInstance("str")]), Ty::Tuple(vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str")]))]
     fn is_subtype_of(from: Ty, to: Ty) {
         let db = setup_db();
         assert!(from.into_type(&db).is_subtype_of(&db, to.into_type(&db)));
@@ -1997,6 +2026,10 @@ mod tests {
     #[test_case(Ty::Union(vec![Ty::IntLiteral(1), Ty::IntLiteral(2)]), Ty::Union(vec![Ty::IntLiteral(1), Ty::IntLiteral(3)]))]
     #[test_case(Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str"))]
     #[test_case(Ty::BuiltinInstance("int"), Ty::IntLiteral(1))]
+    #[test_case(Ty::Tuple(vec![]), Ty::Tuple(vec![Ty::IntLiteral(1)]))]
+    #[test_case(Ty::Tuple(vec![Ty::IntLiteral(42)]), Ty::Tuple(vec![Ty::BuiltinInstance("str")]))]
+    #[test_case(Ty::Tuple(vec![Ty::Todo]), Ty::Tuple(vec![Ty::IntLiteral(2)]))]
+    #[test_case(Ty::Tuple(vec![Ty::IntLiteral(2)]), Ty::Tuple(vec![Ty::Todo]))]
     fn is_not_subtype_of(from: Ty, to: Ty) {
         let db = setup_db();
         assert!(!from.into_type(&db).is_subtype_of(&db, to.into_type(&db)));

crates/red_knot_python_semantic/src/types/infer.rs

@@ -3568,11 +3568,27 @@ impl<'db> TypeInferenceBuilder<'db> {
             ast::Expr::NumberLiteral(_literal) => Type::Todo,
             ast::Expr::BooleanLiteral(_literal) => Type::Todo,
 
-            // TODO: this may be a place we need to revisit with special forms.
             ast::Expr::Subscript(subscript) => {
-                self.infer_subscript_expression(subscript);
+                let ast::ExprSubscript {
+                    value,
+                    slice,
+                    ctx: _,
+                    range: _,
+                } = subscript;
+
+                let value_ty = self.infer_expression(value);
+
+                if value_ty
+                    .into_class_literal_type()
+                    .is_some_and(|class| class.is_known(self.db, KnownClass::Tuple))
+                {
+                    self.infer_tuple_type_expression(slice)
+                } else {
+                    self.infer_type_expression(slice);
+                    // TODO: many other kinds of subscripts
                     Type::Todo
                 }
+            }
 
             ast::Expr::BinOp(binary) => {
                 #[allow(clippy::single_match_else)]
@@ -3591,6 +3607,12 @@ impl<'db> TypeInferenceBuilder<'db> {
                 }
             }
 
+            // TODO PEP 646
+            ast::Expr::Starred(starred) => {
+                self.infer_starred_expression(starred);
+                Type::Todo
+            }
+
             // Forms which are invalid in the context of annotation expressions: we infer their
             // nested expressions as normal expressions, but the type of the top-level expression is
             // always `Type::Unknown` in these cases.
@@ -3667,10 +3689,6 @@ impl<'db> TypeInferenceBuilder<'db> {
                 self.infer_attribute_expression(attribute);
                 Type::Unknown
             }
-            ast::Expr::Starred(starred) => {
-                self.infer_starred_expression(starred);
-                Type::Unknown
-            }
             ast::Expr::List(list) => {
                 self.infer_list_expression(list);
                 Type::Unknown
@@ -3693,6 +3711,59 @@ impl<'db> TypeInferenceBuilder<'db> {
 
         ty
     }
+
+    /// Given the slice of a `tuple[]` annotation, return the type that the annotation represents
+    fn infer_tuple_type_expression(&mut self, tuple_slice: &ast::Expr) -> Type<'db> {
+        /// In most cases, if a subelement of the tuple is inferred as `Todo`,
+        /// we should only infer `Todo` for that specific subelement.
+        /// Certain specific AST nodes can however change the meaning of the entire tuple,
+        /// however: for example, `tuple[int, ...]` or `tuple[int, *tuple[str, ...]]` are a
+        /// homogeneous tuple and a partly homogeneous tuple (respectively) due to the `...`
+        /// and the starred expression (respectively), Neither is supported by us right now,
+        /// so we should infer `Todo` for the *entire* tuple if we encounter one of those elements.
+        /// Even a subscript subelement could alter the type of the entire tuple
+        /// if the subscript is `Unpack[]` (which again, we don't yet support).
+        fn element_could_alter_type_of_whole_tuple(element: &ast::Expr, element_ty: Type) -> bool {
+            element_ty.is_todo()
+                && matches!(
+                    element,
+                    ast::Expr::EllipsisLiteral(_) | ast::Expr::Starred(_) | ast::Expr::Subscript(_)
+                )
+        }
+
+        // TODO:
+        // - homogeneous tuples
+        // - PEP 646
+        match tuple_slice {
+            ast::Expr::Tuple(elements) => {
+                let mut element_types = Vec::with_capacity(elements.len());
+
+                // Whether to infer `Todo` for the whole tuple
+                // (see docstring for `element_could_alter_type_of_whole_tuple`)
+                let mut return_todo = false;
+
+                for element in elements {
+                    let element_ty = self.infer_type_expression(element);
+                    return_todo |= element_could_alter_type_of_whole_tuple(element, element_ty);
+                    element_types.push(element_ty);
+                }
+
+                if return_todo {
+                    Type::Todo
+                } else {
+                    Type::Tuple(TupleType::new(self.db, element_types.into_boxed_slice()))
+                }
+            }
+            single_element => {
+                let single_element_ty = self.infer_type_expression(single_element);
+                if element_could_alter_type_of_whole_tuple(single_element, single_element_ty) {
+                    Type::Todo
+                } else {
+                    Type::Tuple(TupleType::new(self.db, Box::from([single_element_ty])))
+                }
+            }
+        }
+    }
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]