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

Author: ibraheemdev

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

@@ -131,23 +131,23 @@ 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]]`"
+# error: [invalid-assignment] "Object of type `list[Unknown | 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]`"
+# error: [invalid-assignment] "Object of type `list[Unknown | 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]`"
+# error: [invalid-assignment] "Object of type `list[Unknown | 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]`"
+# error: [invalid-assignment] "Object of type `set[Unknown | int | str]` is not assignable to `set[int]`"
 b: set[int] = {1, 2, "3"}
 ```
 
@@ -263,8 +263,7 @@ reveal_type(d)  # revealed: list[int | tuple[int, int]]
 e: list[int] = f(True)
 reveal_type(e)  # revealed: list[int]
 
-# TODO: the RHS should be inferred as `list[Literal["a"]]` here
-# error: [invalid-assignment] "Object of type `list[int | Literal["a"]]` is not assignable to `list[int]`"
+# error: [invalid-assignment] "Object of type `list[Literal["a"]]` is not assignable to `list[int]`"
 g: list[int] = f("a")
 
 # error: [invalid-assignment] "Object of type `list[Literal["a"]]` is not assignable to `tuple[int]`"

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 non-standalone 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 non-standalone 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(expr, TypeContext::default());
+        let _ = self.context.finish();
+        expr_ty
+    }
+
     pub(super) fn finish_expression(mut self) -> ExpressionInference<'db> {
         self.infer_region();