[ty] Infer more precise types for collection literals (#20360)

## Summary

Part of astral-sh/ty#168. Infer more precise types for collection literals (currently, only `list` and `set`). For example,

```py
x = [1, 2, 3] # revealed: list[Unknown | int]
y: list[int] = [1, 2, 3] # revealed: list[int]
```

This could easily be extended to `dict` literals, but I am intentionally limiting scope for now.

Author: ibraheemdev

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

@@ -79,6 +79,78 @@ b: tuple[int] = ("foo",)
 c: tuple[str | int, str] = ([], "foo")
 ```
 
+## Collection literal annotations are understood
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+```py
+import typing
+
+a: list[int] = [1, 2, 3]
+reveal_type(a)  # revealed: list[int]
+
+b: list[int | str] = [1, 2, 3]
+reveal_type(b)  # revealed: list[int | str]
+
+c: typing.List[int] = [1, 2, 3]
+reveal_type(c)  # revealed: list[int]
+
+d: list[typing.Any] = []
+reveal_type(d)  # revealed: list[Any]
+
+e: set[int] = {1, 2, 3}
+reveal_type(e)  # revealed: set[int]
+
+f: set[int | str] = {1, 2, 3}
+reveal_type(f)  # revealed: set[int | str]
+
+g: typing.Set[int] = {1, 2, 3}
+reveal_type(g)  # revealed: set[int]
+
+h: list[list[int]] = [[], [42]]
+reveal_type(h)  # revealed: list[list[int]]
+
+i: list[typing.Any] = [1, 2, "3", ([4],)]
+reveal_type(i)  # revealed: list[Any | int | str | tuple[list[Unknown | int]]]
+
+j: list[tuple[str | int, ...]] = [(1, 2), ("foo", "bar"), ()]
+reveal_type(j)  # revealed: list[tuple[str | int, ...]]
+
+k: list[tuple[list[int], ...]] = [([],), ([1, 2], [3, 4]), ([5], [6], [7])]
+reveal_type(k)  # revealed: list[tuple[list[int], ...]]
+
+l: tuple[list[int], *tuple[list[typing.Any], ...], list[str]] = ([1, 2, 3], [4, 5, 6], [7, 8, 9], ["10", "11", "12"])
+reveal_type(l)  # revealed: tuple[list[int], list[Any | int], list[Any | int], list[str]]
+
+type IntList = list[int]
+
+m: IntList = [1, 2, 3]
+reveal_type(m)  # revealed: list[int]
+
+# TODO: this should type-check and avoid literal promotion
+# error: [invalid-assignment] "Object of type `list[int]` is not assignable to `list[Literal[1, 2, 3]]`"
+n: list[typing.Literal[1, 2, 3]] = [1, 2, 3]
+reveal_type(n)  # revealed: list[Literal[1, 2, 3]]
+
+# TODO: this should type-check and avoid literal promotion
+# error: [invalid-assignment] "Object of type `list[str]` is not assignable to `list[LiteralString]`"
+o: list[typing.LiteralString] = ["a", "b", "c"]
+reveal_type(o)  # revealed: list[LiteralString]
+```
+
+## Incorrect collection literal assignments are complained aobut
+
+```py
+# error: [invalid-assignment] "Object of type `list[int]` is not assignable to `list[str]`"
+a: list[str] = [1, 2, 3]
+
+# error: [invalid-assignment] "Object of type `set[int | str]` is not assignable to `set[int]`"
+b: set[int] = {1, 2, "3"}
+```
+
 ## PEP-604 annotations are supported
 
 ```py

crates/ty_python_semantic/resources/mdtest/del.md

@@ -46,7 +46,7 @@ def delete():
     del d  # error: [unresolved-reference] "Name `d` used when not defined"
 
 delete()
-reveal_type(d)  # revealed: list[@Todo(list literal element type)]
+reveal_type(d)  # revealed: list[Unknown | int]
 
 def delete_element():
     # When the `del` target isn't a name, it doesn't force local resolution.
@@ -62,7 +62,7 @@ def delete_global():
 
 delete_global()
 # Again, the variable should have been removed, but we don't check it.
-reveal_type(d)  # revealed: list[@Todo(list literal element type)]
+reveal_type(d)  # revealed: list[Unknown | int]
 
 def delete_nonlocal():
     e = 2

crates/ty_python_semantic/resources/mdtest/import/dunder_all.md

@@ -783,9 +783,8 @@ class A: ...
 ```py
 from subexporter import *
 
-# TODO: Should be `list[str]`
 # TODO: Should we avoid including `Unknown` for this case?
-reveal_type(__all__)  # revealed: Unknown | list[@Todo(list literal element type)]
+reveal_type(__all__)  # revealed: Unknown | list[Unknown | str]
 
 __all__.append("B")
 

crates/ty_python_semantic/resources/mdtest/literal/collections/list.md

@@ -3,7 +3,33 @@
 ## Empty list
 
 ```py
-reveal_type([])  # revealed: list[@Todo(list literal element type)]
+reveal_type([])  # revealed: list[Unknown]
+```
+
+## List of tuples
+
+```py
+reveal_type([(1, 2), (3, 4)])  # revealed: list[Unknown | tuple[int, int]]
+```
+
+## List of functions
+
+```py
+def a(_: int) -> int:
+    return 0
+
+def b(_: int) -> int:
+    return 1
+
+x = [a, b]
+reveal_type(x)  # revealed: list[Unknown | ((_: int) -> int)]
+```
+
+## Mixed list
+
+```py
+# revealed: list[Unknown | int | tuple[int, int] | tuple[int, int, int]]
+reveal_type([1, (1, 2), (1, 2, 3)])
 ```
 
 ## List comprehensions

crates/ty_python_semantic/resources/mdtest/literal/collections/set.md

@@ -3,7 +3,33 @@
 ## Basic set
 
 ```py
-reveal_type({1, 2})  # revealed: set[@Todo(set literal element type)]
+reveal_type({1, 2})  # revealed: set[Unknown | int]
+```
+
+## Set of tuples
+
+```py
+reveal_type({(1, 2), (3, 4)})  # revealed: set[Unknown | tuple[int, int]]
+```
+
+## Set of functions
+
+```py
+def a(_: int) -> int:
+    return 0
+
+def b(_: int) -> int:
+    return 1
+
+x = {a, b}
+reveal_type(x)  # revealed: set[Unknown | ((_: int) -> int)]
+```
+
+## Mixed set
+
+```py
+# revealed: set[Unknown | int | tuple[int, int] | tuple[int, int, int]]
+reveal_type({1, (1, 2), (1, 2, 3)})
 ```
 
 ## Set comprehensions

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md

@@ -310,26 +310,21 @@ no longer valid in the inner lazy scope.
 def f(l: list[str | None]):
     if l[0] is not None:
         def _():
-            # TODO: should be `str | None`
-            reveal_type(l[0])  # revealed: str | None | @Todo(list literal element type)
-        # TODO: should be of type `list[None]`
+            reveal_type(l[0])  # revealed: str | None | Unknown
         l = [None]
 
 def f(l: list[str | None]):
     l[0] = "a"
     def _():
-        # TODO: should be `str | None`
-        reveal_type(l[0])  # revealed: str | None | @Todo(list literal element type)
-    # TODO: should be of type `list[None]`
+        reveal_type(l[0])  # revealed: str | None | Unknown
     l = [None]
 
 def f(l: list[str | None]):
     l[0] = "a"
     def _():
         l: list[str | None] = [None]
         def _():
-            # TODO: should be `str | None`
-            reveal_type(l[0])  # revealed: @Todo(list literal element type)
+            reveal_type(l[0])  # revealed: str | None
 
     def _():
         def _():

crates/ty_python_semantic/resources/mdtest/subscript/lists.md

@@ -9,13 +9,11 @@ A list can be indexed into with:
 
 ```py
 x = [1, 2, 3]
-reveal_type(x)  # revealed: list[@Todo(list literal element type)]
+reveal_type(x)  # revealed: list[Unknown | int]
 
-# TODO reveal int
-reveal_type(x[0])  # revealed: @Todo(list literal element type)
+reveal_type(x[0])  # revealed: Unknown | int
 
-# TODO reveal list[int]
-reveal_type(x[0:1])  # revealed: list[@Todo(list literal element type)]
+reveal_type(x[0:1])  # revealed: list[Unknown | int]
 
 # error: [invalid-argument-type]
 reveal_type(x["a"])  # revealed: Unknown

crates/ty_python_semantic/resources/mdtest/type_compendium/tuple.md

@@ -55,8 +55,7 @@ def f(x: Iterable[int], y: list[str], z: Never, aa: list[Never], bb: LiskovUncom
 
 reveal_type(tuple((1, 2)))  # revealed: tuple[Literal[1], Literal[2]]
 
-# TODO: should be `tuple[Literal[1], ...]`
-reveal_type(tuple([1]))  # revealed: tuple[@Todo(list literal element type), ...]
+reveal_type(tuple([1]))  # revealed: tuple[Unknown | int, ...]
 
 # error: [invalid-argument-type]
 reveal_type(tuple[int]([1]))  # revealed: tuple[int]

crates/ty_python_semantic/resources/mdtest/unpacking.md

@@ -213,9 +213,8 @@ reveal_type(d)  # revealed: Literal[2]
 
 ```py
 a, b = [1, 2]
-# TODO: should be `int` for both `a` and `b`
-reveal_type(a)  # revealed: @Todo(list literal element type)
-reveal_type(b)  # revealed: @Todo(list literal element type)
+reveal_type(a)  # revealed: Unknown | int
+reveal_type(b)  # revealed: Unknown | int
 ```
 
 ### Simple unpacking

crates/ty_python_semantic/src/types.rs

@@ -1130,11 +1130,30 @@ impl<'db> Type<'db> {
             Type::IntLiteral(_) => Some(KnownClass::Int.to_instance(db)),
             Type::BytesLiteral(_) => Some(KnownClass::Bytes.to_instance(db)),
             Type::ModuleLiteral(_) => Some(KnownClass::ModuleType.to_instance(db)),
+            Type::FunctionLiteral(_) => Some(KnownClass::FunctionType.to_instance(db)),
             Type::EnumLiteral(literal) => Some(literal.enum_class_instance(db)),
             _ => None,
         }
     }
 
+    /// If this type is a literal, promote it to a type that this literal is an instance of.
+    ///
+    /// Note that this function tries to promote literals to a more user-friendly form than their
+    /// fallback instance type. For example, `def _() -> int` is promoted to `Callable[[], int]`,
+    /// as opposed to `FunctionType`.
+    pub(crate) fn literal_promotion_type(self, db: &'db dyn Db) -> Option<Type<'db>> {
+        match self {
+            Type::StringLiteral(_) | Type::LiteralString => Some(KnownClass::Str.to_instance(db)),
+            Type::BooleanLiteral(_) => Some(KnownClass::Bool.to_instance(db)),
+            Type::IntLiteral(_) => Some(KnownClass::Int.to_instance(db)),
+            Type::BytesLiteral(_) => Some(KnownClass::Bytes.to_instance(db)),
+            Type::ModuleLiteral(_) => Some(KnownClass::ModuleType.to_instance(db)),
+            Type::EnumLiteral(literal) => Some(literal.enum_class_instance(db)),
+            Type::FunctionLiteral(literal) => Some(Type::Callable(literal.into_callable_type(db))),
+            _ => None,
+        }
+    }
+
     /// Return a "normalized" version of `self` that ensures that equivalent types have the same Salsa ID.
     ///
     /// A normalized type:
@@ -1704,18 +1723,13 @@ impl<'db> Type<'db> {
                 | Type::IntLiteral(_)
                 | Type::BytesLiteral(_)
                 | Type::ModuleLiteral(_)
-                | Type::EnumLiteral(_),
+                | Type::EnumLiteral(_)
+                | Type::FunctionLiteral(_),
                 _,
             ) => (self.literal_fallback_instance(db)).when_some_and(|instance| {
                 instance.has_relation_to_impl(db, target, relation, visitor)
             }),
 
-            // A `FunctionLiteral` type is a single-valued type like the other literals handled above,
-            // so it also, for now, just delegates to its instance fallback.
-            (Type::FunctionLiteral(_), _) => KnownClass::FunctionType
-                .to_instance(db)
-                .has_relation_to_impl(db, target, relation, visitor),
-
             // The same reasoning applies for these special callable types:
             (Type::BoundMethod(_), _) => KnownClass::MethodType
                 .to_instance(db)
@@ -5979,8 +5993,9 @@ impl<'db> Type<'db> {
                         self
                     }
                 }
-                TypeMapping::PromoteLiterals | TypeMapping::BindLegacyTypevars(_) |
-                TypeMapping::MarkTypeVarsInferable(_) => self,
+                TypeMapping::PromoteLiterals
+                    | TypeMapping::BindLegacyTypevars(_)
+                    | TypeMapping::MarkTypeVarsInferable(_) => self,
                 TypeMapping::Materialize(materialization_kind)  => {
                 Type::TypeVar(bound_typevar.materialize_impl(db, *materialization_kind, visitor))
             }
@@ -6000,10 +6015,10 @@ impl<'db> Type<'db> {
                         self
                     }
                 }
-                TypeMapping::PromoteLiterals |
-                TypeMapping::BindLegacyTypevars(_) |
-                TypeMapping::BindSelf(_) |
-                TypeMapping::ReplaceSelf { .. }
+                TypeMapping::PromoteLiterals
+                    | TypeMapping::BindLegacyTypevars(_)
+                    | TypeMapping::BindSelf(_)
+                    | TypeMapping::ReplaceSelf { .. }
                     => self,
                 TypeMapping::Materialize(materialization_kind)  => Type::NonInferableTypeVar(bound_typevar.materialize_impl(db, *materialization_kind, visitor))
 
@@ -6023,7 +6038,13 @@ impl<'db> Type<'db> {
             }
 
             Type::FunctionLiteral(function) => {
-                Type::FunctionLiteral(function.with_type_mapping(db, type_mapping))
+                let function = Type::FunctionLiteral(function.with_type_mapping(db, type_mapping));
+
+                match type_mapping {
+                    TypeMapping::PromoteLiterals => function.literal_promotion_type(db)
+                        .expect("function literal should have a promotion type"),
+                    _ => function
+                }
             }
 
             Type::BoundMethod(method) => Type::BoundMethod(BoundMethodType::new(
@@ -6129,8 +6150,8 @@ impl<'db> Type<'db> {
                 TypeMapping::ReplaceSelf { .. } |
                 TypeMapping::MarkTypeVarsInferable(_) |
                 TypeMapping::Materialize(_) => self,
-                TypeMapping::PromoteLiterals => self.literal_fallback_instance(db)
-                    .expect("literal type should have fallback instance type"),
+                TypeMapping::PromoteLiterals => self.literal_promotion_type(db)
+                    .expect("literal type should have a promotion type"),
             }
 
             Type::Dynamic(_) => match type_mapping {
@@ -6663,8 +6684,8 @@ pub enum TypeMapping<'a, 'db> {
     Specialization(Specialization<'db>),
     /// Applies a partial specialization to the type
     PartialSpecialization(PartialSpecialization<'a, 'db>),
-    /// Promotes any literal types to their corresponding instance types (e.g. `Literal["string"]`
-    /// to `str`)
+    /// Replaces any literal types with their corresponding promoted type form (e.g. `Literal["string"]`
+    /// to `str`, or `def _() -> int` to `Callable[[], int]`).
     PromoteLiterals,
     /// Binds a legacy typevar with the generic context (class, function, type alias) that it is
     /// being used in.