nits

Author: ibraheemdev

crates/ty_python_semantic/src/types.rs

@@ -1202,16 +1202,24 @@ impl<'db> Type<'db> {
         if yes { self.negate(db) } else { *self }
     }
 
-    /// 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 the type is a union, filters union elements based on the provided predicate.
+    ///
+    /// Otherwise, returns the type unchanged.
+    pub(crate) fn filter_union(
+        self,
+        db: &'db dyn Db,
+        f: impl FnMut(&Type<'db>) -> bool,
+    ) -> Type<'db> {
         if let Type::Union(union) = self {
             union.filter(db, f)
         } else {
             self
         }
     }
 
-    /// Remove the union elements that are not related to `target`.
+    /// If the type is a union, removes union elements that are disjoint from `target`.
+    ///
+    /// Otherwise, returns the type unchanged.
     pub(crate) fn filter_disjoint_elements(
         self,
         db: &'db dyn Db,

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

@@ -2582,28 +2582,25 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
             }
         }
 
-        // Build the specialization once without type context.
+        // Build the specialization once without inferring the type context.
         let isolated_specialization = builder.build(generic_context, *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);
-                }
-            }
+        //
+        // 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.
+        if let Some(call_expression_tcx) = self.call_expression_tcx.annotation
+            && !call_expression_tcx.is_type_var()
+        {
+            // 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.
+            // Build the specialization a second time after inferring the type context.
             let specialization = builder.build(generic_context, *self.call_expression_tcx);
             let return_ty = self.return_ty.apply_specialization(self.db, specialization);