avoid unnecessarily widening generic specializations

Author: ibraheemdev

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

@@ -190,8 +190,7 @@ 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]]
+reveal_type(l)  # revealed: tuple[list[int], list[Any | int], list[Any | int], list[str]]
 
 type IntList = list[int]
 
@@ -416,13 +415,13 @@ a = f("a")
 reveal_type(a)  # revealed: list[Literal["a"]]
 
 b: list[int | Literal["a"]] = f("a")
-reveal_type(b)  # revealed: list[int | Literal["a"]]
+reveal_type(b)  # revealed: list[Literal["a"] | int]
 
 c: list[int | str] = f("a")
-reveal_type(c)  # revealed: list[int | str]
+reveal_type(c)  # revealed: list[str | int]
 
 d: list[int | tuple[int, int]] = f((1, 2))
-reveal_type(d)  # revealed: list[int | tuple[int, int]]
+reveal_type(d)  # revealed: list[tuple[int, int] | int]
 
 e: list[int] = f(True)
 reveal_type(e)  # revealed: list[int]
@@ -437,8 +436,43 @@ def f2[T: int](x: T) -> T:
     return x
 
 i: int = f2(True)
-reveal_type(i)  # revealed: int
+reveal_type(i)  # revealed: Literal[True]
 
 j: int | str = f2(True)
 reveal_type(j)  # revealed: Literal[True]
 ```
+
+Types are not widened unnecessarily:
+
+```py
+def id[T](x: T) -> T:
+    return x
+
+def lst[T](x: T) -> list[T]:
+    return [x]
+
+def _(i: int):
+    a: int | None = i
+    b: int | None = id(i)
+    c: int | str | None = id(i)
+    reveal_type(a)  # revealed: int
+    reveal_type(b)  # revealed: int
+    reveal_type(c)  # revealed: int
+
+    a: list[int | None] | None = [i]
+    b: list[int | None] | None = id([i])
+    c: list[int | None] | int | None = id([i])
+    reveal_type(a)  # revealed: list[int | None]
+
+    # TODO: these should reveal `list[int | None]`
+    # we currently do not use the call expression annotation as type context for argument inference
+    reveal_type(b)  # revealed: list[Unknown | int] | list[int | None] | None
+    reveal_type(c)  # revealed: list[Unknown | int] | list[int | None] | int | None
+
+    a: list[int | None] | None = [i]
+    b: list[int | None] | None = lst(i)
+    c: list[int | None] | int | None = lst(i)
+    reveal_type(a)  # revealed: list[int | None]
+    reveal_type(b)  # revealed: list[int | None]
+    reveal_type(c)  # revealed: list[int | None]
+```

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -907,7 +907,7 @@ grandchild: Node = {"name": "grandchild", "parent": child}
 
 nested: Node = {"name": "n1", "parent": {"name": "n2", "parent": {"name": "n3", "parent": None}}}
 
-# TODO: this should be an error (invalid type for `name` in innermost node)
+# error: [invalid-argument-type] "Invalid argument to key "name" with declared type `str` on TypedDict `Node`: value of type `Literal[3]`"
 nested_invalid: Node = {"name": "n1", "parent": {"name": "n2", "parent": {"name": 3, "parent": None}}}
 ```
 

crates/ty_python_semantic/src/types.rs

@@ -1196,15 +1196,20 @@ 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> {
+    /// Filters union elements based on the provided predicate.
+    pub(crate) fn filter(self, db: &'db dyn Db, f: impl FnMut(&Type<'db>) -> bool) -> Type<'db> {
         if let Type::Union(union) = self {
-            union.filter(db, |elem| !elem.is_disjoint_from(db, target))
+            union.filter(db, f)
         } 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> {
+        self.filter(db, |elem| !elem.is_disjoint_from(db, target))
+    }
+
     /// 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>> {
@@ -10861,9 +10866,9 @@ impl<'db> UnionType<'db> {
     pub(crate) fn filter(
         self,
         db: &'db dyn Db,
-        filter_fn: impl FnMut(&&Type<'db>) -> bool,
+        mut f: impl FnMut(&Type<'db>) -> bool,
     ) -> Type<'db> {
-        Self::from_elements(db, self.elements(db).iter().filter(filter_fn))
+        Self::from_elements(db, self.elements(db).iter().filter(|ty| f(ty)))
     }
 
     pub(crate) fn map_with_boundness(

crates/ty_python_semantic/src/types/call/bind.rs

@@ -2459,19 +2459,22 @@ struct ArgumentTypeChecker<'a, 'db> {
     argument_matches: &'a [MatchedArgument<'db>],
     parameter_tys: &'a mut [Option<Type<'db>>],
     call_expression_tcx: &'a TypeContext<'db>,
+    return_ty: Type<'db>,
     errors: &'a mut Vec<BindingError<'db>>,
 
     specialization: Option<Specialization<'db>>,
 }
 
 impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
+    #[allow(clippy::too_many_arguments)]
     fn new(
         db: &'db dyn Db,
         signature: &'a Signature<'db>,
         arguments: &'a CallArguments<'a, 'db>,
         argument_matches: &'a [MatchedArgument<'db>],
         parameter_tys: &'a mut [Option<Type<'db>>],
         call_expression_tcx: &'a TypeContext<'db>,
+        return_ty: Type<'db>,
         errors: &'a mut Vec<BindingError<'db>>,
     ) -> Self {
         Self {
@@ -2481,6 +2484,7 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
             argument_matches,
             parameter_tys,
             call_expression_tcx,
+            return_ty,
             errors,
             specialization: None,
         }
@@ -2514,31 +2518,12 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
     }
 
     fn infer_specialization(&mut self) {
-        if self.signature.generic_context.is_none() {
+        let Some(gc) = self.signature.generic_context else {
             return;
-        }
+        };
 
         let mut builder = SpecializationBuilder::new(self.db);
 
-        // Note that we infer the annotated type _before_ the arguments if this call is part of
-        // an annotated assignment, to closer match the order of any unions written in the type
-        // annotation.
-        if let Some(return_ty) = self.signature.return_ty
-            && let Some(call_expression_tcx) = self.call_expression_tcx.annotation
-        {
-            match call_expression_tcx {
-                // A type variable is not a useful type-context for expression inference, and applying it
-                // to the return type can lead to confusing unions in nested generic calls.
-                Type::TypeVar(_) => {}
-
-                _ => {
-                    // Ignore any specialization errors here, because the type context is only used as a hint
-                    // to infer a more assignable return type.
-                    let _ = builder.infer(return_ty, call_expression_tcx);
-                }
-            }
-        }
-
         let parameters = self.signature.parameters();
         for (argument_index, adjusted_argument_index, _, argument_type) in
             self.enumerate_argument_types()
@@ -2563,10 +2548,42 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
             }
         }
 
-        self.specialization = self
-            .signature
-            .generic_context
-            .map(|gc| builder.build(gc, *self.call_expression_tcx));
+        // Build the specialization once without type context.
+        let isolated_specialization = builder.build(gc, *self.call_expression_tcx);
+        let isolated_return_ty = self
+            .return_ty
+            .apply_specialization(self.db, isolated_specialization);
+
+        // TODO: Ideally we would infer the annotated type _before_ the arguments if this call is part of an
+        // annotated assignment, to closer match the order of any unions written in the type annotation.
+        if let Some(call_expression_tcx) = self.call_expression_tcx.annotation {
+            match call_expression_tcx {
+                // A type variable is not a useful type-context for expression inference, and applying it
+                // to the return type can lead to confusing unions in nested generic calls.
+                Type::TypeVar(_) => {}
+
+                _ => {
+                    // Ignore any specialization errors here, because the type context is only used as a hint
+                    // to infer a more assignable return type.
+                    let _ = builder.infer(self.return_ty, call_expression_tcx);
+                }
+            }
+
+            // Build the specialization a second time with the type context.
+            let specialization = builder.build(gc, *self.call_expression_tcx);
+            let return_ty = self.return_ty.apply_specialization(self.db, specialization);
+
+            // The type context should only be used to infer a more assignable return type, not an unnecessarily wide
+            // one, so we may ignore the type context here to prefer the narrower return type.
+            if !isolated_return_ty.is_subtype_of(self.db, return_ty) {
+                self.return_ty = return_ty;
+                self.specialization = Some(specialization);
+                return;
+            }
+        }
+
+        self.return_ty = isolated_return_ty;
+        self.specialization = Some(isolated_specialization);
     }
 
     fn check_argument_type(
@@ -2754,8 +2771,8 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
         }
     }
 
-    fn finish(self) -> Option<Specialization<'db>> {
-        self.specialization
+    fn finish(self) -> (Option<Specialization<'db>>, Type<'db>) {
+        (self.specialization, self.return_ty)
     }
 }
 
@@ -2908,6 +2925,7 @@ impl<'db> Binding<'db> {
             &self.argument_matches,
             &mut self.parameter_tys,
             call_expression_tcx,
+            self.return_ty,
             &mut self.errors,
         );
 
@@ -2916,10 +2934,7 @@ impl<'db> Binding<'db> {
         checker.infer_specialization();
 
         checker.check_argument_types();
-        self.specialization = checker.finish();
-        if let Some(specialization) = self.specialization {
-            self.return_ty = self.return_ty.apply_specialization(db, specialization);
-        }
+        (self.specialization, self.return_ty) = checker.finish();
     }
 
     pub(crate) fn set_return_type(&mut self, return_ty: Type<'db>) {

crates/ty_python_semantic/src/types/generics.rs

@@ -1198,6 +1198,7 @@ impl<'db> SpecializationBuilder<'db> {
                     let tcx = tcx_specialization.and_then(|specialization| {
                         specialization.get(self.db, variable.bound_typevar)
                     });
+
                     ty = ty.map(|ty| ty.promote_literals(self.db, TypeContext::new(tcx)));
                 }
 
@@ -1220,7 +1221,7 @@ impl<'db> SpecializationBuilder<'db> {
     pub(crate) fn infer(
         &mut self,
         formal: Type<'db>,
-        mut actual: Type<'db>,
+        actual: Type<'db>,
     ) -> Result<(), SpecializationError<'db>> {
         if formal == actual {
             return Ok(());
@@ -1249,9 +1250,11 @@ impl<'db> SpecializationBuilder<'db> {
             return Ok(());
         }
 
-        // 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);
+        // Remove the union elements that are not related to `formal`.
+        //
+        // For example, if `formal` is `list[T]` and `actual` is `list[int] | None`, we want to specialize `T`
+        // to `int`.
+        let actual = actual.filter_disjoint_elements(self.db, formal);
 
         match (formal, actual) {
             // TODO: We haven't implemented a full unification solver yet. If typevars appear in

crates/ty_python_semantic/src/types/infer.rs

@@ -391,7 +391,7 @@ impl<'db> TypeContext<'db> {
             .and_then(|ty| ty.known_specialization(db, known_class))
     }
 
-    pub(crate) fn map_annotation(self, f: impl FnOnce(Type<'db>) -> Type<'db>) -> Self {
+    pub(crate) fn map(self, f: impl FnOnce(Type<'db>) -> Type<'db>) -> Self {
         Self {
             annotation: self.annotation.map(f),
         }

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

@@ -5819,6 +5819,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             parenthesized: _,
         } = tuple;
 
+        // Remove any union elements of that are unrelated to the tuple type.
+        let tcx = tcx.map(|annotation| {
+            annotation.filter_disjoint_elements(self.db(), KnownClass::Tuple.to_instance(self.db()))
+        });
+
         let annotated_tuple = tcx
             .known_specialization(self.db(), KnownClass::Tuple)
             .and_then(|specialization| {
@@ -5884,7 +5889,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         } = dict;
 
         // Validate `TypedDict` dictionary literal assignments.
-        if let Some(typed_dict) = tcx.annotation.and_then(Type::as_typed_dict)
+        if let Some(tcx) = tcx.annotation
+            && let Some(typed_dict) = tcx.filter(self.db(), Type::is_typed_dict).as_typed_dict()
             && let Some(ty) = self.infer_typed_dict_expression(dict, typed_dict)
         {
             return ty;
@@ -5964,9 +5970,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             return None;
         };
 
-        let tcx = tcx.map_annotation(|annotation| {
-            // Remove any union elements of `annotation` that are not related to `collection_ty`.
-            // e.g. `annotation: list[int] | None => list[int]` if `collection_ty: list`
+        // Remove any union elements of that are unrelated to the collection type.
+        //
+        // For example, we only want the `list[int]` from `annotation: list[int] | None` if
+        // `collection_ty` is `list`.
+        let tcx = tcx.map(|annotation| {
             let collection_ty = collection_class.to_instance(self.db());
             annotation.filter_disjoint_elements(self.db(), collection_ty)
         });