re-infer RHS of annotated assignments in isolation for assignability diagnostics

crates/ty_python_semantic/src/types.rs

@@ -25,7 +25,8 @@ pub use self::diagnostic::TypeCheckDiagnostics;
 pub(crate) use self::diagnostic::register_lints;
 pub(crate) use self::infer::{
     TypeContext, infer_deferred_types, infer_definition_types, infer_expression_type,
-    infer_expression_types, infer_scope_types, static_expression_truthiness,
+    infer_expression_types, infer_isolated_expression, infer_scope_types,
+    static_expression_truthiness,
 };
 pub(crate) use self::signatures::{CallableSignature, Signature};
 pub(crate) use self::subclass_of::{SubclassOfInner, SubclassOfType};

crates/ty_python_semantic/src/types/diagnostic.rs

@@ -6,7 +6,7 @@ use super::{
     add_inferred_python_version_hint_to_diagnostic,
 };
 use crate::lint::{Level, LintRegistryBuilder, LintStatus};
-use crate::semantic_index::definition::Definition;
+use crate::semantic_index::definition::{Definition, DefinitionKind};
 use crate::semantic_index::place::{PlaceTable, ScopedPlaceId};
 use crate::suppression::FileSuppressionId;
 use crate::types::call::CallError;
@@ -19,7 +19,7 @@ use crate::types::string_annotation::{
 };
 use crate::types::{
     ClassType, DynamicType, LintDiagnosticGuard, Protocol, ProtocolInstanceType, SubclassOfInner,
-    binding_type,
+    binding_type, infer_isolated_expression,
 };
 use crate::types::{SpecialFormType, Type, protocol_class::ProtocolClass};
 use crate::util::diagnostics::format_enumeration;
@@ -1940,15 +1940,24 @@ fn report_invalid_assignment_with_message(
     }
 }
 
-pub(super) fn report_invalid_assignment(
-    context: &InferContext,
+pub(super) fn report_invalid_assignment<'db>(
+    context: &InferContext<'db, '_>,
     node: AnyNodeRef,
+    definition: Definition<'db>,
     target_ty: Type,
-    source_ty: Type,
+    mut source_ty: Type<'db>,
 ) {
     let settings =
         DisplaySettings::from_possibly_ambiguous_type_pair(context.db(), target_ty, source_ty);
 
+    if let DefinitionKind::AnnotatedAssignment(annotated_assignment) = definition.kind(context.db())
+        && let Some(value) = annotated_assignment.value(context.module())
+    {
+        // Re-infer the RHS of the annotated assignment, ignoring the type context, for more precise
+        // error messages.
+        source_ty = infer_isolated_expression(context.db(), definition.scope(context.db()), value);
+    }
+
     report_invalid_assignment_with_message(
         context,
         node,

crates/ty_python_semantic/src/types/infer.rs

@@ -37,6 +37,7 @@
 //! be considered a bug.)
 
 use ruff_db::parsed::{ParsedModuleRef, parsed_module};
+use ruff_python_ast as ast;
 use ruff_text_size::Ranged;
 use rustc_hash::FxHashMap;
 use salsa;
@@ -217,6 +218,24 @@ fn infer_expression_types_impl<'db>(
     .finish_expression()
 }
 
+/// Infer the type of an expression in isolation.
+///
+/// The type returned by this function may be different than the type of the expression
+/// if it was inferred within its region, as it does not account for surrounding type context.
+/// This can be useful to re-infer the type of an expression for diagnostics.
+pub(crate) fn infer_isolated_expression<'db>(
+    db: &'db dyn Db,
+    scope: ScopeId<'db>,
+    expr: &ast::Expr,
+) -> Type<'db> {
+    let file = scope.file(db);
+    let module = parsed_module(db, file).load(db);
+    let index = semantic_index(db, file);
+
+    TypeInferenceBuilder::new(db, InferenceRegion::Scope(scope), index, &module)
+        .infer_isolated_expression(expr)
+}
+
 fn expression_cycle_recover<'db>(
     db: &'db dyn Db,
     _value: &ExpressionInference<'db>,

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

@@ -1522,7 +1522,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         }
 
         if !bound_ty.is_assignable_to(db, declared_ty) {
-            report_invalid_assignment(&self.context, node, declared_ty, bound_ty);
+            report_invalid_assignment(&self.context, node, binding, declared_ty, bound_ty);
             // allow declarations to override inference in case of invalid assignment
             bound_ty = declared_ty;
         }
@@ -1679,9 +1679,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     report_invalid_assignment(
                         &self.context,
                         node,
+                        definition,
                         declared_ty.inner_type(),
                         inferred_ty,
                     );
+
                     // if the assignment is invalid, fall back to assuming the annotation is correct
                     (declared_ty, declared_ty.inner_type())
                 }
@@ -5336,29 +5338,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             panic!("Typeshed should always have a `{name}` class in `builtins.pyi` with a single type variable")
         });
 
-        let mut elements_are_assignable = true;
-        let mut inferred_elt_tys = Vec::with_capacity(elts.len());
-
-        // Infer the type of each element in the collection literal.
-        for elt in elts {
-            let inferred_elt_ty = self.infer_expression(elt, TypeContext::new(annotated_elts_ty));
-            inferred_elt_tys.push(inferred_elt_ty);
-
-            if let Some(annotated_elts_ty) = annotated_elts_ty {
-                elements_are_assignable &=
-                    inferred_elt_ty.is_assignable_to(self.db(), annotated_elts_ty);
-            }
-        }
-
         // Create a set of constraints to infer a precise type for `T`.
         let mut builder = SpecializationBuilder::new(self.db());
 
         match annotated_elts_ty {
-            // If the inferred type of any element is not assignable to the type annotation, we
-            // ignore it, as to provide a more precise error message.
-            Some(_) if !elements_are_assignable => {}
-
-            // Otherwise, the annotated type acts as a constraint for `T`.
+            // The annotated type acts as a constraint for `T`.
             //
             // Note that we infer the annotated type _before_ the elements, to closer match the order
             // of any unions written in the type annotation.
@@ -5372,7 +5356,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         }
 
         // The inferred type of each element acts as an additional constraint on `T`.
-        for inferred_elt_ty in inferred_elt_tys {
+        for elt in elts {
+            let inferred_elt_ty = self.infer_expression(elt, TypeContext::new(annotated_elts_ty));
+
             // Convert any element literals to their promoted type form to avoid excessively large
             // unions for large nested list literals, which the constraint solver struggles with.
             let inferred_elt_ty =
@@ -9032,6 +9018,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         }
     }
 
+    /// Infer the type of the given expression in isolation, ignoring the surrounding region.
+    pub(super) fn infer_isolated_expression(mut self, expr: &ast::Expr) -> Type<'db> {
+        let expr_ty = self.infer_expression_impl(expr, TypeContext::default());
+        let _ = self.context.finish();
+        expr_ty
+    }
+
     pub(super) fn finish_expression(mut self) -> ExpressionInference<'db> {
         self.infer_region();