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}