menva_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, Data, DeriveInput, Fields, FieldsNamed, LitStr};
4
5#[proc_macro_derive(FromEnv, attributes(env_prefix))]
6pub fn from_env_macro(input: TokenStream) -> TokenStream {
7    let input = parse_macro_input!(input as DeriveInput);
8    let named_fields = get_named_fields(&input);
9    let from_env_with_prefix_fn = generate_from_env_with_prefix(named_fields);
10    let from_env_fn = generate_from_env(named_fields);
11
12    let struct_name = input.ident;
13
14    let expanded = quote! {
15        impl #struct_name {
16            #from_env_with_prefix_fn
17            #from_env_fn
18        }
19    };
20
21    TokenStream::from(expanded)
22}
23
24fn get_named_fields(input: &DeriveInput) -> &FieldsNamed {
25    if let Data::Struct(data) = &input.data {
26        if let Fields::Named(named_fields) = &data.fields {
27            named_fields
28        } else {
29            panic!("FromEnvWithPrefix can only be derived for structs with named fields");
30        }
31    } else {
32        panic!("FromEnvWithPrefix can only be derived for structs");
33    }
34}
35
36fn generate_from_env_with_prefix(named_fields: &FieldsNamed) -> proc_macro2::TokenStream {
37    let fields = named_fields.named.iter().map(|field| {
38        let field_name = field.ident.as_ref().unwrap();
39        let field_name_str = field_name.to_string().to_uppercase();
40
41        quote! {
42            #field_name: {
43                let env_var_name = format!("{}{}", prefix, #field_name_str);
44                std::env::var(&env_var_name)
45                    .expect(&format!("Environment variable `{}` not set", env_var_name))
46                    .parse()
47                    .expect(&format!("Failed to parse `{}`", env_var_name))
48            }
49        }
50    }).collect::<Vec<_>>();
51
52    quote! {
53        pub fn from_env_with_prefix(prefix: &str) -> Self {
54            Self {
55                #(#fields),*
56            }
57        }
58    }
59}
60
61fn generate_from_env(named_fields: &FieldsNamed) -> proc_macro2::TokenStream {
62    let fields = named_fields
63        .named
64        .iter()
65        .map(|field| {
66            let field_name = field.ident.as_ref().unwrap();
67            let field_name_str = field_name.to_string().to_uppercase();
68
69            let env_var_name = get_env_var_name(field, field_name_str);
70
71            quote! {
72            #field_name:  std::env::var(#env_var_name)
73            .expect(&format!("Environment variable `{}` not set", #env_var_name))
74            .parse()
75            .expect(&format!("Failed to parse `{}`", #env_var_name))
76            }
77        })
78        .collect::<Vec<_>>();
79
80    quote! {
81        pub fn from_env() -> Self {
82            Self {
83                #(#fields),*
84            }
85        }
86    }
87}
88
89fn get_env_var_name(field: &syn::Field, field_name_str: String) -> String {
90    field
91        .attrs
92        .iter()
93        .find(|attr| attr.path().is_ident("env_prefix"))
94        .map(|attr| {
95            let mut prefix = attr
96                .parse_args::<LitStr>()
97                .unwrap_or_else(|_| {
98                    panic!("Prefix for `{}` has a problem", field_name_str)
99                })
100                .value();
101            prefix.push_str(&field_name_str);
102            prefix
103        })
104        .unwrap_or(field_name_str)
105}