From 2a2b5d7b6274b80321a96031af18969445d78651 Mon Sep 17 00:00:00 2001
From: Benoit Jeaurond <benoit@jeaurond.dev>
Date: Sun, 16 Jan 2022 16:16:11 -0500
Subject: [PATCH] feat: add serde derives to enums, if specified

---
 sea-orm-codegen/src/entity/active_enum.rs |  8 ++++++--
 sea-orm-codegen/src/entity/writer.rs      | 18 ++++++++----------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/sea-orm-codegen/src/entity/active_enum.rs b/sea-orm-codegen/src/entity/active_enum.rs
index d92423f25..800da59e5 100644
--- a/sea-orm-codegen/src/entity/active_enum.rs
+++ b/sea-orm-codegen/src/entity/active_enum.rs
@@ -1,3 +1,4 @@
+use crate::WithSerde;
 use heck::CamelCase;
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
@@ -9,7 +10,7 @@ pub struct ActiveEnum {
 }
 
 impl ActiveEnum {
-    pub fn impl_active_enum(&self) -> TokenStream {
+    pub fn impl_active_enum(&self, with_serde: &WithSerde) -> TokenStream {
         let enum_name = &self.enum_name;
         let enum_iden = format_ident!("{}", enum_name.to_camel_case());
         let values = &self.values;
@@ -17,8 +18,11 @@ impl ActiveEnum {
             .values
             .iter()
             .map(|v| format_ident!("{}", v.to_camel_case()));
+
+        let extra_derive = with_serde.extra_derive();
+
         quote! {
-            #[derive(Debug, Clone, PartialEq, EnumIter, DeriveActiveEnum)]
+            #[derive(Debug, Clone, PartialEq, EnumIter, DeriveActiveEnum #extra_derive)]
             #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = #enum_name)]
             pub enum #enum_iden {
                 #(
diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs
index 60a01eeb2..326a6e668 100644
--- a/sea-orm-codegen/src/entity/writer.rs
+++ b/sea-orm-codegen/src/entity/writer.rs
@@ -81,25 +81,25 @@ impl FromStr for WithSerde {
 impl EntityWriter {
     pub fn generate(self, expanded_format: bool, with_serde: WithSerde) -> WriterOutput {
         let mut files = Vec::new();
-        files.extend(self.write_entities(expanded_format, with_serde));
+        files.extend(self.write_entities(expanded_format, &with_serde));
         files.push(self.write_mod());
         files.push(self.write_prelude());
         if !self.enums.is_empty() {
-            files.push(self.write_sea_orm_active_enums());
+            files.push(self.write_sea_orm_active_enums(&with_serde));
         }
         WriterOutput { files }
     }
 
-    pub fn write_entities(&self, expanded_format: bool, with_serde: WithSerde) -> Vec<OutputFile> {
+    pub fn write_entities(&self, expanded_format: bool, with_serde: &WithSerde) -> Vec<OutputFile> {
         self.entities
             .iter()
             .map(|entity| {
                 let mut lines = Vec::new();
                 Self::write_doc_comment(&mut lines);
                 let code_blocks = if expanded_format {
-                    Self::gen_expanded_code_blocks(entity, &with_serde)
+                    Self::gen_expanded_code_blocks(entity, with_serde)
                 } else {
-                    Self::gen_compact_code_blocks(entity, &with_serde)
+                    Self::gen_compact_code_blocks(entity, with_serde)
                 };
                 Self::write(&mut lines, code_blocks);
                 OutputFile {
@@ -147,20 +147,18 @@ impl EntityWriter {
         }
     }
 
-    pub fn write_sea_orm_active_enums(&self) -> OutputFile {
+    pub fn write_sea_orm_active_enums(&self, with_serde: &WithSerde) -> OutputFile {
         let mut lines = Vec::new();
         Self::write_doc_comment(&mut lines);
         Self::write(
             &mut lines,
-            vec![quote! {
-                use sea_orm::entity::prelude::*;
-            }],
+            vec![Self::gen_import(with_serde)],
         );
         lines.push("".to_owned());
         let code_blocks = self
             .enums
             .iter()
-            .map(|(_, active_enum)| active_enum.impl_active_enum())
+            .map(|(_, active_enum)| active_enum.impl_active_enum(with_serde))
             .collect();
         Self::write(&mut lines, code_blocks);
         OutputFile {