From 7cf319556d17c57cb95a91b96b573468a310bc44 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Schl=C3=BCter=20Lukas?= Date: Sun, 11 Jun 2023 19:09:58 +0200 Subject: [PATCH 1/2] fix: Only generate `use` statements once per enum --- sea-orm-codegen/src/entity/writer.rs | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs index 36da98910..42597fce3 100644 --- a/sea-orm-codegen/src/entity/writer.rs +++ b/sea-orm-codegen/src/entity/writer.rs @@ -460,15 +460,18 @@ impl EntityWriter { entity .columns .iter() - .fold(TokenStream::new(), |mut ts, col| { + .fold((TokenStream::new(), Vec::new()), |(mut ts, mut enums), col| { if let sea_query::ColumnType::Enum { name, .. } = &col.col_type { - let enum_name = format_ident!("{}", name.to_string().to_upper_camel_case()); - ts.extend([quote! { - use super::sea_orm_active_enums::#enum_name; - }]); + if !enums.contains(&name) { + enums.push(name); + let enum_name = format_ident!("{}", name.to_string().to_upper_camel_case()); + ts.extend([quote! { + use super::sea_orm_active_enums::#enum_name; + }]); + } } - ts - }) + (ts, enums) + }).0 } pub fn gen_model_struct( From 9ec81077a984311973faac3084caa041e42ef44f Mon Sep 17 00:00:00 2001 From: Zander Milroy Date: Thu, 2 Nov 2023 00:51:33 +0000 Subject: [PATCH 2/2] Adds gen_import_active_enum tests --- sea-orm-codegen/src/entity/writer.rs | 128 ++++++++++++++++++++++++++- 1 file changed, 127 insertions(+), 1 deletion(-) diff --git a/sea-orm-codegen/src/entity/writer.rs b/sea-orm-codegen/src/entity/writer.rs index 9eadeae0d..6600f4624 100644 --- a/sea-orm-codegen/src/entity/writer.rs +++ b/sea-orm-codegen/src/entity/writer.rs @@ -822,7 +822,8 @@ mod tests { }; use pretty_assertions::assert_eq; use proc_macro2::TokenStream; - use sea_query::{ColumnType, ForeignKeyAction, RcOrArc}; + use quote::quote; + use sea_query::{Alias, ColumnType, ForeignKeyAction, RcOrArc, SeaRc}; use std::io::{self, BufRead, BufReader, Read}; fn setup() -> Vec { @@ -2283,4 +2284,129 @@ mod tests { Ok(()) } + + #[test] + fn test_gen_import_active_enum() -> io::Result<()> { + let entities = vec![ + Entity { + table_name: "tea_pairing".to_owned(), + columns: vec![ + Column { + name: "id".to_owned(), + col_type: ColumnType::Integer, + auto_increment: true, + not_null: true, + unique: false, + }, + Column { + name: "first_tea".to_owned(), + col_type: ColumnType::Enum { + name: SeaRc::new(Alias::new("tea_enum")), + variants: vec![ + SeaRc::new(Alias::new("everyday_tea")), + SeaRc::new(Alias::new("breakfast_tea")), + ], + }, + auto_increment: false, + not_null: true, + unique: false, + }, + Column { + name: "second_tea".to_owned(), + col_type: ColumnType::Enum { + name: SeaRc::new(Alias::new("tea_enum")), + variants: vec![ + SeaRc::new(Alias::new("everyday_tea")), + SeaRc::new(Alias::new("breakfast_tea")), + ], + }, + auto_increment: false, + not_null: true, + unique: false, + }, + ], + relations: vec![], + conjunct_relations: vec![], + primary_keys: vec![PrimaryKey { + name: "id".to_owned(), + }], + }, + Entity { + table_name: "tea_pairing_with_size".to_owned(), + columns: vec![ + Column { + name: "id".to_owned(), + col_type: ColumnType::Integer, + auto_increment: true, + not_null: true, + unique: false, + }, + Column { + name: "first_tea".to_owned(), + col_type: ColumnType::Enum { + name: SeaRc::new(Alias::new("tea_enum")), + variants: vec![ + SeaRc::new(Alias::new("everyday_tea")), + SeaRc::new(Alias::new("breakfast_tea")), + ], + }, + auto_increment: false, + not_null: true, + unique: false, + }, + Column { + name: "second_tea".to_owned(), + col_type: ColumnType::Enum { + name: SeaRc::new(Alias::new("tea_enum")), + variants: vec![ + SeaRc::new(Alias::new("everyday_tea")), + SeaRc::new(Alias::new("breakfast_tea")), + ], + }, + auto_increment: false, + not_null: true, + unique: false, + }, + Column { + name: "size".to_owned(), + col_type: ColumnType::Enum { + name: SeaRc::new(Alias::new("tea_size")), + variants: vec![ + SeaRc::new(Alias::new("small")), + SeaRc::new(Alias::new("medium")), + SeaRc::new(Alias::new("huge")), + ], + }, + auto_increment: false, + not_null: true, + unique: false, + }, + ], + relations: vec![], + conjunct_relations: vec![], + primary_keys: vec![PrimaryKey { + name: "id".to_owned(), + }], + }, + ]; + + assert_eq!( + quote!( + use super::sea_orm_active_enums::TeaEnum; + ) + .to_string(), + EntityWriter::gen_import_active_enum(&entities[0]).to_string() + ); + + assert_eq!( + quote!( + use super::sea_orm_active_enums::TeaEnum; + use super::sea_orm_active_enums::TeaSize; + ) + .to_string(), + EntityWriter::gen_import_active_enum(&entities[1]).to_string() + ); + + Ok(()) + } }