Skip to main content

dotenv_derive/
dotenv_derive.rs

1use proc_macro::TokenStream;
2use proc_macro2::Ident;
3use quote::quote;
4use syn::{
5    Data, DeriveInput, Expr, Fields, LitStr, Token,
6    parse::{Parse, ParseStream},
7    parse_macro_input,
8};
9
10#[proc_macro_derive(FromEnv, attributes(env))]
11pub fn derive_from_env(input: TokenStream) -> TokenStream {
12    let input = parse_macro_input!(input as DeriveInput);
13    let name = &input.ident;
14
15    let fields = match &input.data {
16        Data::Struct(data) => match &data.fields {
17            Fields::Named(fields) => &fields.named,
18            _ => panic!("FromEnv only supports structs with named fields"),
19        },
20        _ => panic!("FromEnv only supports structs"),
21    };
22
23    let field_assignments: Vec<_> = fields
24        .iter()
25        .map(|f| {
26            let field_name = &f.ident;
27            let field_name_str = field_name.as_ref().unwrap().to_string();
28            let screaming = to_screaming_snake(&field_name_str);
29            let ty = &f.ty;
30
31            let attrs: Vec<EnvAttr> = f
32                .attrs
33                .iter()
34                .filter(|a| a.path().is_ident("env"))
35                .map(|a| a.parse_args::<EnvAttr>().unwrap())
36                .collect();
37
38            let attr = attrs.first();
39
40            match attr {
41                // #[env(rename = "VAR")] — explicit name, FromEnvValue parsing
42                Some(EnvAttr {
43                    rename: Some(var_name),
44                    default: None,
45                    with: None,
46                }) => {
47                    let var = var_name.value();
48                    quote! {
49                        #field_name: {
50                            let val = ::std::env::var(#var)
51                                .map_err(|_| ::dotenv::FromEnvError::missing(#var))?;
52                            <#ty as ::dotenv::FromEnvValue>::from_env_value(val.clone())
53                                .map_err(|e| ::dotenv::FromEnvError::invalid(#var, val, e))?
54                        }
55                    }
56                }
57                // #[env(rename = "VAR", default)] — explicit name, fallback to Default::default()
58                Some(EnvAttr {
59                    rename: Some(var_name),
60                    default: Some(DefaultKind::Standard),
61                    with: None,
62                }) => {
63                    let var = var_name.value();
64                    quote! {
65                        #field_name: {
66                            match ::std::env::var(#var) {
67                                Ok(val) => <#ty as ::dotenv::FromEnvValue>::from_env_value(val.clone())
68                                    .map_err(|e| ::dotenv::FromEnvError::invalid(#var, val, e))?,
69                                Err(_) => ::std::default::Default::default(),
70                            }
71                        }
72                    }
73                }
74                // #[env(rename = "VAR", default = EXPR)] — explicit name, fallback to expr
75                Some(EnvAttr {
76                    rename: Some(var_name),
77                    default: Some(DefaultKind::Expr(expr)),
78                    with: None,
79                }) => {
80                    let var = var_name.value();
81                    quote! {
82                        #field_name: {
83                            match ::std::env::var(#var) {
84                                Ok(val) => <#ty as ::dotenv::FromEnvValue>::from_env_value(val.clone())
85                                    .map_err(|e| ::dotenv::FromEnvError::invalid(#var, val, e))?,
86                                Err(_) => #expr,
87                            }
88                        }
89                    }
90                }
91                // #[env(rename = "VAR", with = "func")] — explicit name, custom parser
92                Some(EnvAttr {
93                    rename: Some(var_name),
94                    default: None,
95                    with: Some(func),
96                }) => {
97                    let var = var_name.value();
98                    let func = Ident::new(&func.value(), proc_macro2::Span::call_site());
99                    quote! {
100                        #field_name: {
101                            let val = ::std::env::var(#var)
102                                .map_err(|_| ::dotenv::FromEnvError::missing(#var))?;
103                            #func(#var, &val)?
104                        }
105                    }
106                }
107                // #[env(with = "func")] — default naming, custom parser
108                Some(EnvAttr {
109                    rename: None,
110                    default: None,
111                    with: Some(func),
112                }) => {
113                    let func = Ident::new(&func.value(), proc_macro2::Span::call_site());
114                    quote! {
115                        #field_name: {
116                            let var_name = ::std::format!("{}{}", prefix, #screaming);
117                            let val = ::std::env::var(&var_name)
118                                .map_err(|_| ::dotenv::FromEnvError::missing(var_name.clone()))?;
119                            #func(&var_name, &val)?
120                        }
121                    }
122                }
123                // #[env(default)] — unconditional Default::default()
124                // Does NOT read the env var; the default is used unconditionally.
125                // For reading + fallback, combine with rename: #[env(rename = "VAR", default)]
126                Some(EnvAttr {
127                    rename: None,
128                    default: Some(DefaultKind::Standard),
129                    with: None,
130                }) => {
131                    quote! {
132                        #field_name: ::std::default::Default::default()
133                    }
134                }
135                // #[env(default = EXPR)] — unconditional expression
136                // Does NOT read the env var; the expression is used unconditionally.
137                Some(EnvAttr {
138                    rename: None,
139                    default: Some(DefaultKind::Expr(expr)),
140                    with: None,
141                }) => {
142                    quote! {
143                        #field_name: #expr
144                    }
145                }
146                // #[env()] or #[env] with no args — same as no attribute
147                Some(EnvAttr {
148                    rename: None,
149                    default: None,
150                    with: None,
151                }) => {
152                    quote! {
153                        #field_name: <#ty as ::dotenv::FromEnvAuto>::from_env_auto(
154                            &::std::format!("{}{}_", prefix, #screaming),
155                            &::std::format!("{}{}", prefix, #screaming),
156                        )?
157                    }
158                }
159                // No attribute → auto-dispatch via FromEnvAuto
160                //   - If the type implements FromEnv (nested struct): calls from_env_with_prefix
161                //   - Otherwise (leaf type like String, u32): reads env var + FromStr
162                None => {
163                    quote! {
164                        #field_name: <#ty as ::dotenv::FromEnvAuto>::from_env_auto(
165                            &::std::format!("{}{}_", prefix, #screaming),
166                            &::std::format!("{}{}", prefix, #screaming),
167                        )?
168                    }
169                }
170                // Invalid combination of attributes
171                _ => {
172                    panic!(
173                        "invalid `#[env(...)]` attributes on field `{}`: supported are \
174                     `#[env(rename = \"...\")]`, `#[env(rename = \"...\", with = \"...\")]`, \
175                     `#[env(with = \"...\")]`, `#[env(default)]`, \
176                     `#[env(default = ...)]`, or no attribute",
177                        field_name_str
178                    );
179                }
180            }
181        })
182        .collect();
183
184    let expanded = quote! {
185        #[automatically_derived]
186        impl ::dotenv::FromEnv for #name {
187            fn from_env_with_prefix(prefix: &str) -> ::std::result::Result<Self, ::dotenv::FromEnvError> {
188                Ok(Self {
189                    #(#field_assignments,)*
190                })
191            }
192        }
193    };
194
195    TokenStream::from(expanded)
196}
197
198fn to_screaming_snake(s: &str) -> String {
199    let mut result = String::new();
200    let chars: Vec<char> = s.chars().collect();
201    let mut i = 0;
202    while i < chars.len() {
203        if i > 0 && chars[i].is_ascii_uppercase() {
204            let prev_lower = chars[i - 1].is_ascii_lowercase();
205            let next_lower = i + 1 < chars.len() && chars[i + 1].is_ascii_lowercase();
206            if prev_lower || (next_lower && i > 0 && chars[i - 1].is_ascii_uppercase()) {
207                result.push('_');
208            }
209        }
210        result.push(chars[i].to_ascii_uppercase());
211        i += 1;
212    }
213    result
214}
215
216struct EnvAttr {
217    rename: Option<LitStr>,
218    default: Option<DefaultKind>,
219    with: Option<LitStr>,
220}
221
222enum DefaultKind {
223    Standard,
224    Expr(Expr),
225}
226
227impl Parse for EnvAttr {
228    fn parse(input: ParseStream) -> syn::Result<Self> {
229        let mut rename = None;
230        let mut default = None;
231        let mut with = None;
232
233        while !input.is_empty() {
234            let ident: syn::Ident = input.parse()?;
235            match ident.to_string().as_str() {
236                "rename" => {
237                    input.parse::<Token![=]>()?;
238                    rename = Some(input.parse()?);
239                }
240                "default" => {
241                    if input.peek(Token![=]) {
242                        input.parse::<Token![=]>()?;
243                        default = Some(DefaultKind::Expr(input.parse()?));
244                    } else {
245                        default = Some(DefaultKind::Standard);
246                    }
247                }
248                "with" => {
249                    input.parse::<Token![=]>()?;
250                    with = Some(input.parse()?);
251                }
252                other => {
253                    return Err(syn::Error::new(ident.span(), format!("unknown env attribute: `{other}`")));
254                }
255            }
256            if !input.is_empty() {
257                input.parse::<Token![,]>()?;
258            }
259        }
260
261        Ok(EnvAttr {
262            rename,
263            default,
264            with,
265        })
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    #[test]
274    fn test_to_screaming_snake() {
275        assert_eq!(to_screaming_snake("url"), "URL");
276        assert_eq!(to_screaming_snake("my_field"), "MY_FIELD");
277        assert_eq!(to_screaming_snake("myField"), "MY_FIELD");
278        assert_eq!(to_screaming_snake("XMLParser"), "XML_PARSER");
279        assert_eq!(to_screaming_snake("database_url"), "DATABASE_URL");
280        assert_eq!(to_screaming_snake("db"), "DB");
281        assert_eq!(to_screaming_snake("a"), "A");
282        assert_eq!(to_screaming_snake(""), "");
283    }
284}