Merge pull request #4272 from Ten0/auto_type_binary_operators

Auto type binary operators
diesel-rs · Sep 27, 2024 · ce6fa0b · ce6fa0b
2 parents a020eb5 + cfd2a56
commit ce6fa0b
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 0 deletions.
diff --git a/diesel/src/lib.rs b/diesel/src/lib.rs
@@ -651,6 +651,11 @@ pub mod helper_types {
     #[allow(non_camel_case_types)] // required for `#[auto_type]`
     pub type insert_into<T> = crate::query_builder::IncompleteInsertStatement<T>;
 
+    /// Represents the return type of [`diesel::update`]
+    #[allow(non_camel_case_types)] // required for `#[auto_type]`
+    pub type update<T> =
+        UpdateStatement<<T as HasTable>::Table, <T as IntoUpdateTarget>::WhereClause>;
+
     /// Represents the return type of [`diesel::insert_or_ignore_into`]
     #[allow(non_camel_case_types)] // required for `#[auto_type]`
     pub type insert_or_ignore_into<T> = crate::query_builder::IncompleteInsertOrIgnoreStatement<T>;

diff --git a/diesel_derives/tests/auto_type.rs b/diesel_derives/tests/auto_type.rs
@@ -481,6 +481,15 @@ fn delete_returning() -> _ {
     delete(users::table).returning(users::id)
 }
 
+#[cfg(feature = "postgres")]
+#[auto_type]
+fn update_and_binary_operator_and_block() -> _ {
+    update(pg_extras::table).set(pg_extras::timestamp.eq(pg_extras::timestamp + {
+        let v: diesel::data_types::PgInterval = 1.year();
+        v
+    }))
+}
+
 // #[auto_type]
 // fn test_sql_fragment() -> _ {
 //     sql("foo")

diff --git a/dsl_auto_type/src/auto_type/expression_type_inference.rs b/dsl_auto_type/src/auto_type/expression_type_inference.rs
@@ -302,6 +302,54 @@ impl TypeInferrer<'_> {
                     }
                 }
             }
+            (syn::Expr::Binary(binary_expression), None) => {
+                let op_span = Span::mixed_site().located_at(binary_expression.op.span());
+                let trait_name = match binary_expression.op {
+                    syn::BinOp::Add(_) => "Add",
+                    syn::BinOp::Sub(_) => "Sub",
+                    syn::BinOp::Mul(_) => "Mul",
+                    syn::BinOp::Div(_) => "Div",
+                    syn::BinOp::Rem(_) => "Rem",
+                    syn::BinOp::BitXor(_) => "BitXor",
+                    syn::BinOp::BitAnd(_) => "BitAnd",
+                    syn::BinOp::BitOr(_) => "BitOr",
+                    syn::BinOp::Shl(_) => "Shl",
+                    syn::BinOp::Shr(_) => "Shr",
+                    syn::BinOp::And(_)
+                    | syn::BinOp::Or(_)
+                    | syn::BinOp::Eq(_)
+                    | syn::BinOp::Lt(_)
+                    | syn::BinOp::Le(_)
+                    | syn::BinOp::Ne(_)
+                    | syn::BinOp::Ge(_)
+                    | syn::BinOp::Gt(_) => return Ok(parse_quote!(bool)),
+                    syn::BinOp::AddAssign(_)
+                    | syn::BinOp::SubAssign(_)
+                    | syn::BinOp::MulAssign(_)
+                    | syn::BinOp::DivAssign(_)
+                    | syn::BinOp::RemAssign(_)
+                    | syn::BinOp::BitXorAssign(_)
+                    | syn::BinOp::BitAndAssign(_)
+                    | syn::BinOp::BitOrAssign(_)
+                    | syn::BinOp::ShlAssign(_)
+                    | syn::BinOp::ShrAssign(_) => return Ok(parse_quote!(())),
+                    _ => {
+                        // This is here because the `BinOp` enum is marked as #[non_exhaustive],
+                        // but in effect we really support all the variants
+                        return Err(syn::Error::new(
+                            op_span,
+                            format_args!(
+                                "unsupported binary operator for auto_type: {:?}",
+                                binary_expression.op
+                            ),
+                        ));
+                    }
+                };
+                let trait_name_ident = syn::Ident::new(trait_name, op_span);
+                let left_type = self.infer_expression_type(&binary_expression.left, None);
+                let right_type = self.infer_expression_type(&binary_expression.right, None);
+                parse_quote!(<#left_type as ::core::ops::#trait_name_ident<#right_type>>::Output)
+            }
             (_, None) => {
                 return Err(syn::Error::new(
                     expr.span(),