Implement rename attribute for enums

julianskartor · Sep 5, 2024 · 2f74642 · 2f74642
1 parent c57814f
commit 2f74642
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/econf-derive/src/lib.rs b/econf-derive/src/lib.rs
@@ -4,7 +4,7 @@ use proc_macro::TokenStream;
 
 use proc_macro2::{Ident, TokenStream as TokenStream2};
 use quote::quote;
-use syn::{parse_macro_input, Data, DeriveInput, Field, Fields, LitStr};
+use syn::{parse_macro_input, Attribute, Data, DeriveInput, Field, Fields, LitStr, Variant};
 
 #[proc_macro_derive(LoadEnv, attributes(econf))]
 pub fn load_env(input: TokenStream) -> TokenStream {
@@ -37,9 +37,9 @@ fn is_skip(f: &Field) -> bool {
     })
 }
 
-fn find_renaming(f: &Field) -> Option<String> {
+fn find_renaming(attrs: &[Attribute]) -> Option<String> {
     let mut rename = None;
-    for attr in &f.attrs {
+    for attr in attrs {
         if attr.path().is_ident("econf") {
             attr.parse_nested_meta(|meta| {
                 if meta.path.is_ident("rename") {
@@ -56,6 +56,14 @@ fn find_renaming(f: &Field) -> Option<String> {
     rename
 }
 
+fn find_field_renaming(f: &Field) -> Option<String> {
+    find_renaming(&f.attrs)
+}
+
+fn find_variant_renaming(v: &Variant) -> Option<String> {
+    find_renaming(&v.attrs)
+}
+
 fn content(name: &Ident, data: &Data) -> TokenStream2 {
     match data {
         Data::Struct(data) => match &data.fields {
@@ -67,7 +75,7 @@ fn content(name: &Ident, data: &Data) -> TokenStream2 {
                             #ident: self.#ident,
                         };
                     }
-                    match find_renaming(f) {
+                    match find_field_renaming(f) {
                         Some(overwritten_name) => quote! {
                             #ident: self.#ident.load(&(path.to_owned() + "_" + #overwritten_name), loader),
                         },
@@ -89,7 +97,7 @@ fn content(name: &Ident, data: &Data) -> TokenStream2 {
                     if is_skip(f) {
                         return quote! { self.#i, };
                     }
-                    match find_renaming(f) {
+                    match find_field_renaming(f) {
                         Some(overwritten_name) => quote! {
                             self.#i.load(&(path.to_owned() + "_" + #overwritten_name), loader),
                         },
@@ -114,7 +122,11 @@ fn content(name: &Ident, data: &Data) -> TokenStream2 {
             });
 
             let enums0 = data.variants.iter().map(|_| &name);
-            let enums1 = data.variants.iter().map(|f| &f.ident);
+            let enums1 = data.variants.iter().map(|f| {
+                find_variant_renaming(f)
+                    .map(|overwritten_name| Ident::new(&overwritten_name, f.ident.span()))
+                    .unwrap_or_else(|| f.ident.clone())
+            });
             let enums2 = data.variants.iter().map(|f| &f.ident);
 
             quote! {

diff --git a/econf/tests/basics.rs b/econf/tests/basics.rs
@@ -1,4 +1,5 @@
 use econf::LoadEnv;
+use log::*;
 use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet, LinkedList, VecDeque};
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
 use std::path::PathBuf;
@@ -361,33 +362,44 @@ struct NestedRenamed {
     s: String,
 }
 
+#[derive(Debug, Eq, PartialEq, LoadEnv)]
+enum EnumRenamed {
+    V1,
+    #[econf(rename = "example_3")]
+    V2,
+}
+
 #[derive(LoadEnv)]
 struct Renamed {
     v1: bool,
     #[econf(rename = "example_1")]
     v2: u32,
     #[econf(rename = "example_2")]
     v3: NestedRenamed,
+    v4: EnumRenamed,
 }
 
 #[test]
 fn renamed() {
     std::env::set_var("RENAMED_V1", "true");
     std::env::set_var("RENAMED_EXAMPLE_1", "42");
     std::env::set_var("RENAMED_EXAMPLE_2_S", "renamed text");
+    std::env::set_var("RENAMED_V4", "example_3");
 
     let a = Renamed {
         v1: false,
         v2: 0,
         v3: NestedRenamed {
             s: "initial".to_string(),
         },
+        v4: EnumRenamed::V1,
     };
 
     let a = econf::load(a, "renamed");
     assert_eq!(a.v1, true);
     assert_eq!(a.v2, 42);
     assert_eq!(a.v3.s, "renamed text".to_string());
+    assert_eq!(a.v4, EnumRenamed::V2);
 }
 
 #[derive(LoadEnv)]