Check SQLAlchemy update values against mapped fields

pyrefly/lib/alt/expr.rs

@@ -687,6 +687,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 self.check_pytorch_tensor_cuda_call(x, &callee_ty, errors);
                 self.check_pytorch_print_tensor(x, &callee_ty, errors);
                 self.check_pytorch_redundant_to_call(x, &callee_ty, errors);
+                self.check_sqlalchemy_update_values_call(x, errors);
                 if let Some(d) = self.call_to_dict(&callee_ty, &x.arguments) {
                     self.dict_infer(&d, hint, x.range, errors)
                 } else if let Some(ty) = self
@@ -1043,6 +1044,131 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             .emit();
     }
 
+    /// Validate `sqlalchemy.update(Model).values(field=...)` keyword arguments against the
+    /// `Mapped[...]` fields declared on `Model`.
+    fn check_sqlalchemy_update_values_call(&self, x: &ExprCall, errors: &ErrorCollector) {
+        let Expr::Attribute(attr_expr) = &*x.func else {
+            return;
+        };
+        if attr_expr.attr.id.as_str() != "values" {
+            return;
+        }
+        let Some(model) = self.sqlalchemy_update_model_from_chain(&attr_expr.value) else {
+            return;
+        };
+        let mapped_fields = self.sqlalchemy_mapped_model_fields(&model);
+        if mapped_fields.is_empty() {
+            return;
+        }
+        for keyword in &x.arguments.keywords {
+            let Some(field_identifier) = keyword.arg.as_ref() else {
+                continue;
+            };
+            let field_name = &field_identifier.id;
+            if mapped_fields.contains_key(field_name) {
+                let expected_ty = self.sqlalchemy_mapped_field_type(&model, field_name);
+                let got_ty = self.expr_infer(&keyword.value, errors);
+                if !self.is_subset_eq(&got_ty, &expected_ty) {
+                    self.error(
+                        errors,
+                        keyword.value.range(),
+                        ErrorKind::BadArgumentType,
+                        format!(
+                            "`{}` is not assignable to field `{field_name}` with type `{}`",
+                            self.for_display(got_ty),
+                            self.for_display(expected_ty)
+                        ),
+                    );
+                }
+            } else {
+                let mut builder = errors.error_builder(
+                    field_identifier.range(),
+                    ErrorKind::BadArgumentType,
+                    format!("Unexpected SQLAlchemy update field `{field_name}`"),
+                );
+                if let Some(suggestion) = best_suggestion(
+                    field_name,
+                    mapped_fields.keys().map(|candidate| (candidate, 0usize)),
+                ) {
+                    builder = builder.with_detail(format!("Did you mean `{suggestion}`?"));
+                }
+                builder.emit();
+            }
+        }
+    }
+
+    fn sqlalchemy_update_model_from_chain(&self, expr: &Expr) -> Option<Class> {
+        let Expr::Call(call) = expr else {
+            return None;
+        };
+        if self.is_sqlalchemy_update_call(call) {
+            return self.sqlalchemy_update_model_from_call(call);
+        }
+        let Expr::Attribute(attr_expr) = &*call.func else {
+            return None;
+        };
+        self.sqlalchemy_update_model_from_chain(&attr_expr.value)
+    }
+
+    fn is_sqlalchemy_update_call(&self, call: &ExprCall) -> bool {
+        let errors = self.error_swallower();
+        let func_ty = self.expr_infer(&call.func, &errors);
+        let Some(FunctionKind::Def(func_id)) = func_ty.to_func_kind() else {
+            return false;
+        };
+        func_id.name.as_str() == "update"
+            && func_id.module.name().as_str().starts_with("sqlalchemy")
+    }
+
+    fn sqlalchemy_update_model_from_call(&self, call: &ExprCall) -> Option<Class> {
+        let model_expr = call.arguments.args.first()?;
+        let errors = self.error_swallower();
+        match self.expr_infer(model_expr, &errors) {
+            Type::ClassDef(cls) => Some(cls),
+            Type::Type(boxed) => match *boxed {
+                Type::ClassType(cls) => Some(cls.into_class_object()),
+                _ => None,
+            },
+            _ => None,
+        }
+    }
+
+    fn sqlalchemy_mapped_model_fields(&self, model: &Class) -> SmallMap<Name, ()> {
+        self.get_class_field_map(model)
+            .into_iter()
+            .filter_map(|(name, field)| {
+                if name.as_str().starts_with('_') || field.is_class_var() {
+                    return None;
+                }
+                let (_, annotation, _) = field.for_variance_inference();
+                if annotation.is_some_and(|annotation| {
+                    Self::is_sqlalchemy_mapped_annotation(annotation.get_type())
+                }) {
+                    Some((name, ()))
+                } else {
+                    None
+                }
+            })
+            .collect()
+    }
+
+    fn is_sqlalchemy_mapped_annotation(ty: &Type) -> bool {
+        match ty {
+            Type::ClassType(cls) => {
+                cls.has_qname("sqlalchemy.orm", "Mapped")
+                    || cls.has_qname("sqlalchemy.orm.base", "Mapped")
+            }
+            Type::Annotated(inner, _) => Self::is_sqlalchemy_mapped_annotation(inner),
+            _ => false,
+        }
+    }
+
+    fn sqlalchemy_mapped_field_type(&self, model: &Class, field_name: &Name) -> Type {
+        let model_ty = self.instantiate(model);
+        let errors = self.error_swallower();
+        self.attr_infer_for_type(&model_ty, field_name, TextRange::default(), &errors, None)
+    }
+
     fn tuple_infer(&self, x: &ExprTuple, hint: Option<HintRef>, errors: &ErrorCollector) -> Type {
         let owner = Owner::new();
         let (hint_ts, default_hint) = if let Some(hint) = &hint {

pyrefly/lib/test/descriptors.rs

@@ -702,6 +702,14 @@ class Mapped[T]:
     def __delete__(self, instance) -> None: ...
     "#,
     );
+    env.add(
+        "sqlalchemy.sql.dml",
+        r#"
+class Update:
+    def where(self, *criteria: object) -> Update: ...
+    def values(self, **kwargs: object) -> Update: ...
+    "#,
+    );
     env.add_with_path(
         "sqlalchemy.orm.decl_api",
         "sqlalchemy/orm/decl_api.py",
@@ -715,7 +723,14 @@ from .base import Mapped as Mapped
 from .decl_api import DeclarativeBase as DeclarativeBase
     "#,
     );
-    env.add_with_path("sqlalchemy", "sqlalchemy/__init__.py", "");
+    env.add_with_path(
+        "sqlalchemy",
+        "sqlalchemy/__init__.py",
+        r#"
+from .sql.dml import Update as Update
+def update(table: object) -> Update: ...
+    "#,
+    );
     env
 }
 
@@ -760,6 +775,26 @@ class User(Base):
     "#,
 );
 
+testcase!(
+    test_sqlalchemy_update_values_checks_mapped_fields,
+    sqlalchemy_mapped_env(),
+    r#"
+import sqlalchemy as sa
+from sqlalchemy.orm import DeclarativeBase, Mapped
+
+class Base(DeclarativeBase):
+    pass
+
+class User(Base):
+    id: Mapped[int]
+    name: Mapped[str]
+
+sa.update(User).where(User.id == 1).values(name="alice", id=1)
+sa.update(User).where(User.id == 1).values(name=0)  # E: `Literal[0]` is not assignable to field `name` with type `str`
+sa.update(User).where(User.id == 1).values(nam="alice")  # E: Unexpected SQLAlchemy update field `nam`
+    "#,
+);
+
 testcase!(
     test_stub_annotation_only_descriptor_has_descriptor_semantics,
     stub_descriptor_env(),