improve bidirectional inference in infer_collection_literal

Author: mtshiba

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

@@ -150,7 +150,81 @@ r: dict[int | str, int | str] = {1: 1, 2: 2, 3: 3}
 reveal_type(r)  # revealed: dict[int | str, int | str]
 ```
 
-## Incorrect collection literal assignments are complained aobut
+## Optional collection literal annotations are understood
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+```py
+import typing
+
+a: list[int] | None = [1, 2, 3]
+reveal_type(a)  # revealed: list[int]
+
+b: list[int | str] | None = [1, 2, 3]
+reveal_type(b)  # revealed: list[int | str]
+
+c: typing.List[int] | None = [1, 2, 3]
+reveal_type(c)  # revealed: list[int]
+
+d: list[typing.Any] | None = []
+reveal_type(d)  # revealed: list[Any]
+
+e: set[int] | None = {1, 2, 3}
+reveal_type(e)  # revealed: set[int]
+
+f: set[int | str] | None = {1, 2, 3}
+reveal_type(f)  # revealed: set[int | str]
+
+g: typing.Set[int] | None = {1, 2, 3}
+reveal_type(g)  # revealed: set[int]
+
+h: list[list[int]] | None = [[], [42]]
+reveal_type(h)  # revealed: list[list[int]]
+
+i: list[typing.Any] | None = [1, 2, "3", ([4],)]
+reveal_type(i)  # revealed: list[Any | int | str | tuple[list[Unknown | int]]]
+
+j: list[tuple[str | int, ...]] | None = [(1, 2), ("foo", "bar"), ()]
+reveal_type(j)  # revealed: list[tuple[str | int, ...]]
+
+k: list[tuple[list[int], ...]] | None = [([],), ([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]] | None = ([1, 2, 3], [4, 5, 6], [7, 8, 9], ["10", "11", "12"])
+# TODO: this should be `tuple[list[int], list[Any | int], list[Any | int], list[str]]`
+reveal_type(l)  # revealed: tuple[list[Unknown | int], list[Unknown | int], list[Unknown | int], list[Unknown | str]]
+
+type IntList = list[int]
+
+m: IntList | None = [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[Unknown | int]` is not assignable to `list[Literal[1, 2, 3]] | None`"
+n: list[typing.Literal[1, 2, 3]] | None = [1, 2, 3]
+# TODO: this should be `list[Literal[1, 2, 3]]` at this scope
+reveal_type(n)  # revealed: list[Literal[1, 2, 3]] | None
+
+# TODO: this should type-check and avoid literal promotion
+# error: [invalid-assignment] "Object of type `list[Unknown | str]` is not assignable to `list[LiteralString] | None`"
+o: list[typing.LiteralString] | None = ["a", "b", "c"]
+# TODO: this should be `list[LiteralString]` at this scope
+reveal_type(o)  # revealed: list[LiteralString] | None
+
+p: dict[int, int] | None = {}
+reveal_type(p)  # revealed: dict[int, int]
+
+q: dict[int | str, int] | None = {1: 1, 2: 2, 3: 3}
+reveal_type(q)  # revealed: dict[int | str, int]
+
+r: dict[int | str, int | str] | None = {1: 1, 2: 2, 3: 3}
+reveal_type(r)  # revealed: dict[int | str, int | str]
+```
+
+## Incorrect collection literal assignments are complained about
 
 ```py
 # error: [invalid-assignment] "Object of type `list[Unknown | int]` is not assignable to `list[str]`"

crates/ty_python_semantic/src/types.rs

@@ -1140,6 +1140,15 @@ impl<'db> Type<'db> {
         if yes { self.negate(db) } else { *self }
     }
 
+    /// Remove the union elements that are not related to `target`.
+    pub(crate) fn filter_disjoint_elements(self, db: &'db dyn Db, target: Type<'db>) -> Type<'db> {
+        if let Type::Union(union) = self {
+            union.filter(db, |elem| !elem.is_disjoint_from(db, target))
+        } else {
+            self
+        }
+    }
+
     /// Returns the fallback instance type that a literal is an instance of, or `None` if the type
     /// is not a literal.
     pub(crate) fn literal_fallback_instance(self, db: &'db dyn Db) -> Option<Type<'db>> {

crates/ty_python_semantic/src/types/generics.rs

@@ -1159,13 +1159,9 @@ impl<'db> SpecializationBuilder<'db> {
             return Ok(());
         }
 
-        if let Type::Union(union) = actual {
-            // For example, if `formal` is `list[T]` and `actual` is `list[int] | None`, we want to specialize `T` to `int`.
-            // So, here we remove the union elements that are not related to `formal`.
-            actual = union.filter(self.db, |actual_elem| {
-                !actual_elem.is_disjoint_from(self.db, formal)
-            });
-        }
+        // For example, if `formal` is `list[T]` and `actual` is `list[int] | None`, we want to specialize `T` to `int`.
+        // So, here we remove the union elements that are not related to `formal`.
+        actual = actual.filter_disjoint_elements(self.db, formal);
 
         match (formal, actual) {
             (Type::Union(_), Type::Union(_)) => {

crates/ty_python_semantic/src/types/infer.rs

@@ -390,6 +390,12 @@ impl<'db> TypeContext<'db> {
         self.annotation
             .and_then(|ty| ty.known_specialization(known_class, db))
     }
+
+    pub(crate) fn map_annotation(self, f: impl FnOnce(Type<'db>) -> Type<'db>) -> Self {
+        Self {
+            annotation: self.annotation.map(f),
+        }
+    }
 }
 
 /// Returns the statically-known truthiness of a given expression.

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -5560,6 +5560,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             panic!("Typeshed should always have a `{name}` class in `builtins.pyi`")
         });
 
+        let tcx = tcx.map_annotation(|annotation| {
+            // For example, if `collection_ty` is `list` and `annotation` is `list[int] | None`,
+            // remove any union elements of `annotation` that are not related to `collection_ty`.
+            let collection_ty = collection_class.to_instance(self.db());
+            annotation.filter_disjoint_elements(self.db(), collection_ty)
+        });
+
         // Extract the annotated type of `T`, if provided.
         let annotated_elt_tys = tcx
             .known_specialization(collection_class, self.db())