diff --git a/pyrefly/lib/alt/expr.rs b/pyrefly/lib/alt/expr.rs index c3386e44fc..bdcefe4f30 100644 --- a/pyrefly/lib/alt/expr.rs +++ b/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 { + 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 { + 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 { + 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, errors: &ErrorCollector) -> Type { let owner = Owner::new(); let (hint_ts, default_hint) = if let Some(hint) = &hint { diff --git a/pyrefly/lib/test/descriptors.rs b/pyrefly/lib/test/descriptors.rs index ed567018fd..cda7c5908c 100644 --- a/pyrefly/lib/test/descriptors.rs +++ b/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(),