Refactoring

SeaQL · Jul 11, 2023 · 89095d7 · 89095d7
1 parent 2284bb3
commit 89095d7
Showing 1 changed file with 19 additions and 16 deletions.
diff --git a/sea-query-derive/src/lib.rs b/sea-query-derive/src/lib.rs
@@ -46,15 +46,11 @@ fn impl_iden_for_unit_struct(
     ident: &proc_macro2::Ident,
     table_name: &str,
 ) -> proc_macro2::TokenStream {
-    let flagged_path = if cfg!(feature = "sea-orm") {
-        quote!(sea_orm::sea_query)
-    } else {
-        quote!(sea_query)
-    };
+    let sea_query_path = sea_query_path();
 
     let prepare = if must_be_valid_iden(table_name) {
         quote! {
-            fn prepare(&self, s: &mut dyn ::std::fmt::Write, q: #flagged_path::Quote) {
+            fn prepare(&self, s: &mut dyn ::std::fmt::Write, q: #sea_query_path::Quote) {
                 write!(s, "{}", q.left()).unwrap();
                 self.unquoted(s);
                 write!(s, "{}", q.right()).unwrap();
@@ -65,7 +61,7 @@ fn impl_iden_for_unit_struct(
     };
 
     quote! {
-        impl #flagged_path::Iden for #ident {
+        impl #sea_query_path::Iden for #ident {
             #prepare
 
             fn unquoted(&self, s: &mut dyn ::std::fmt::Write) {
@@ -83,6 +79,8 @@ fn impl_iden_for_enum<'a, T>(
 where
     T: Iterator<Item = &'a Variant>,
 {
+    let sea_query_path = sea_query_path();
+
     let mut is_all_valid = true;
 
     let match_arms = match variants
@@ -98,14 +96,9 @@ where
         Err(e) => return e.to_compile_error(),
     };
 
-    let flagged_path = if cfg!(feature = "sea-orm") {
-        quote!(sea_orm::sea_query)
-    } else {
-        quote!(sea_query)
-    };
     let prepare = if is_all_valid {
         quote! {
-            fn prepare(&self, s: &mut dyn ::std::fmt::Write, q: #flagged_path::Quote) {
+            fn prepare(&self, s: &mut dyn ::std::fmt::Write, q: #sea_query_path::Quote) {
                 write!(s, "{}", q.left()).unwrap();
                 self.unquoted(s);
                 write!(s, "{}", q.right()).unwrap();
@@ -116,7 +109,7 @@ where
     };
 
     quote! {
-        impl #flagged_path::Iden for #ident {
+        impl #sea_query_path::Iden for #ident {
             #prepare
 
             fn unquoted(&self, s: &mut dyn ::std::fmt::Write) {
@@ -163,6 +156,8 @@ pub fn derive_iden(input: TokenStream) -> TokenStream {
 
 #[proc_macro_derive(IdenStatic, attributes(iden, method))]
 pub fn derive_iden_static(input: TokenStream) -> TokenStream {
+    let sea_query_path = sea_query_path();
+
     let DeriveInput {
         ident, data, attrs, ..
     } = parse_macro_input!(input);
@@ -185,7 +180,7 @@ pub fn derive_iden_static(input: TokenStream) -> TokenStream {
                 return quote! {
                     #impl_iden
 
-                    impl sea_query::IdenStatic for #ident {
+                    impl #sea_query_path::IdenStatic for #ident {
                         fn as_str(&self) -> &'static str {
                             #table_name
                         }
@@ -224,7 +219,7 @@ pub fn derive_iden_static(input: TokenStream) -> TokenStream {
     let output = quote! {
         #impl_iden
 
-        impl sea_query::IdenStatic for #ident {
+        impl #sea_query_path::IdenStatic for #ident {
             fn as_str(&self) -> &'static str {
                 match self {
                     #match_arms
@@ -241,3 +236,11 @@ pub fn derive_iden_static(input: TokenStream) -> TokenStream {
 
     output.into()
 }
+
+fn sea_query_path() -> proc_macro2::TokenStream {
+    if cfg!(feature = "sea-orm") {
+        quote!(sea_orm::sea_query)
+    } else {
+        quote!(sea_query)
+    }
+}