Skip to main content

dotenv_derive/
dotenv_derive.rs

1use proc_macro::TokenStream;
2use proc_macro2::TokenStream as TokenStream2;
3use quote::quote;
4use syn::{
5    Data, DeriveInput, Expr, Fields, GenericArgument, LitStr, Path, PathArguments, Token, Type,
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    match expand(&input) {
14        Ok(tokens) => tokens.into(),
15        // Emit proper `compile_error!`s instead of panicking.
16        Err(err) => err.to_compile_error().into(),
17    }
18}
19
20fn expand(input: &DeriveInput) -> syn::Result<TokenStream2> {
21    let name = &input.ident;
22
23    let fields = match &input.data {
24        Data::Struct(data) => match &data.fields {
25            Fields::Named(fields) => &fields.named,
26            _ => {
27                return Err(syn::Error::new_spanned(
28                    input,
29                    "FromEnv only supports structs with named fields",
30                ));
31            }
32        },
33        _ => {
34            return Err(syn::Error::new_spanned(input, "FromEnv only supports structs"));
35        }
36    };
37
38    let mut field_assignments = Vec::new();
39    for f in fields {
40        let Some(field_name) = f.ident.as_ref() else {
41            continue;
42        };
43        let field_name_str = field_name.to_string();
44        let screaming = to_screaming_snake(&field_name_str);
45        let ty = &f.ty;
46
47        let mut attrs = Vec::new();
48        for a in f.attrs.iter().filter(|a| a.path().is_ident("env")) {
49            attrs.push(a.parse_args::<EnvAttr>()?);
50        }
51        if attrs.len() > 1 {
52            return Err(syn::Error::new(
53                field_name.span(),
54                "multiple `#[env(...)]` attributes on one field; combine them into a single attribute",
55            ));
56        }
57        let attr = attrs.first();
58
59        let expr = field_expr(field_name, &screaming, ty, attr)?;
60        field_assignments.push(quote! { #field_name: #expr });
61    }
62
63    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
64    Ok(quote! {
65        #[automatically_derived]
66        impl #impl_generics ::dotenv::FromEnv for #name #ty_generics #where_clause {
67            fn from_env_with_prefix(prefix: &str) -> ::std::result::Result<Self, ::dotenv::FromEnvError> {
68                Ok(Self {
69                    #(#field_assignments,)*
70                })
71            }
72        }
73    })
74}
75
76/// Build the value expression for a single field.
77fn field_expr(
78    field_name: &syn::Ident,
79    screaming: &str,
80    ty: &Type,
81    attr: Option<&EnvAttr>,
82) -> syn::Result<TokenStream2> {
83    let inner = option_inner(ty);
84
85    let (rename, default, with) = match attr {
86        Some(EnvAttr {
87            rename,
88            default,
89            with,
90        }) => {
91            if with.is_some() && default.is_some() {
92                return Err(syn::Error::new(field_name.span(), "`with` cannot be combined with `default`"));
93            }
94            (rename.as_ref(), default.as_ref(), with.as_ref())
95        }
96        None => (None, None, None),
97    };
98
99    // `#[env(rename = "VAR", with = "func")]`
100    if let (Some(var), None, Some(path)) = (rename, default, with) {
101        return Ok(quote! {{
102            let val = ::std::env::var(#var)
103                .map_err(|_| ::dotenv::FromEnvError::missing(#var))?;
104            #path(#var, &val)?
105        }});
106    }
107
108    // `#[env(with = "func")]`
109    if let (None, None, Some(path)) = (rename, default, with) {
110        return Ok(quote! {{
111            let var_name = ::std::format!("{}{}", prefix, #screaming);
112            let val = ::std::env::var(&var_name)
113                .map_err(|_| ::dotenv::FromEnvError::missing(var_name.clone()))?;
114            #path(&var_name, &val)?
115        }});
116    }
117
118    if let Some(inner) = inner {
119        // `Option<T>` fields are optional.
120        return Ok(match (rename, default) {
121            (Some(var), None) => quote! {
122                match ::std::env::var(#var) {
123                    ::std::result::Result::Ok(val) if !val.is_empty() => ::std::option::Option::Some(
124                        <#inner as ::dotenv::FromEnvValue>::from_env_value(&val)
125                            .map_err(|e| ::dotenv::FromEnvError::invalid(#var, val, e))?
126                    ),
127                    _ => ::std::option::Option::None,
128                }
129            },
130            (Some(var), Some(kind)) => {
131                let fallback = fallback_tokens(kind);
132                quote! {
133                    match ::std::env::var(#var) {
134                        ::std::result::Result::Ok(val) if !val.is_empty() => ::std::option::Option::Some(
135                            <#inner as ::dotenv::FromEnvValue>::from_env_value(&val)
136                                .map_err(|e| ::dotenv::FromEnvError::invalid(#var, val, e))?
137                        ),
138                        _ => #fallback,
139                    }
140                }
141            }
142            (None, Some(kind)) => {
143                let fallback = fallback_tokens(kind);
144                quote! {{
145                    let var_name = ::std::format!("{}{}", prefix, #screaming);
146                    match ::std::env::var(&var_name) {
147                        ::std::result::Result::Ok(val) if !val.is_empty() => ::std::option::Option::Some(
148                            <#inner as ::dotenv::FromEnvValue>::from_env_value(&val)
149                                .map_err(|e| ::dotenv::FromEnvError::invalid(var_name.as_str(), val, e))?
150                        ),
151                        _ => #fallback,
152                    }
153                }}
154            }
155            (None, None) => quote! {{
156                let var_name = ::std::format!("{}{}", prefix, #screaming);
157                match ::std::env::var(&var_name) {
158                    ::std::result::Result::Ok(val) if !val.is_empty() => ::std::option::Option::Some(
159                        <#inner as ::dotenv::FromEnvValue>::from_env_value(&val)
160                            .map_err(|e| ::dotenv::FromEnvError::invalid(var_name.as_str(), val, e))?
161                    ),
162                    _ => ::std::option::Option::None,
163                }
164            }},
165        });
166    }
167
168    Ok(match (rename, default) {
169        // `#[env(rename = "VAR")]`
170        (Some(var), None) => quote! {{
171            let val = ::std::env::var(#var)
172                .map_err(|_| ::dotenv::FromEnvError::missing(#var))?;
173            <#ty as ::dotenv::FromEnvValue>::from_env_value(&val)
174                .map_err(|e| ::dotenv::FromEnvError::invalid(#var, val, e))?
175        }},
176        // `#[env(rename = "VAR", default)]` / `#[env(rename = "VAR", default = expr)]`
177        (Some(var), Some(kind)) => {
178            let fallback = fallback_tokens(kind);
179            quote! {
180                match ::std::env::var(#var) {
181                    ::std::result::Result::Ok(val) => <#ty as ::dotenv::FromEnvValue>::from_env_value(&val)
182                        .map_err(|e| ::dotenv::FromEnvError::invalid(#var, val, e))?,
183                    ::std::result::Result::Err(_) => #fallback,
184                }
185            }
186        }
187        // `#[env(default)]` / `#[env(default = expr)]` (no rename):
188        // read the derived variable name first, then fall back.
189        (None, Some(kind)) => {
190            let fallback = fallback_tokens(kind);
191            quote! {
192                match <#ty as ::dotenv::FromEnvAuto>::from_env_auto(
193                    &::std::format!("{}{}_", prefix, #screaming),
194                    &::std::format!("{}{}", prefix, #screaming),
195                ) {
196                    ::std::result::Result::Ok(v) => v,
197                    ::std::result::Result::Err(::dotenv::FromEnvError::Missing(_)) => #fallback,
198                    ::std::result::Result::Err(e) => return ::std::result::Result::Err(e),
199                }
200            }
201        }
202        // No attribute (or `#[env()]`): auto-dispatch.
203        (None, None) => quote! {
204            <#ty as ::dotenv::FromEnvAuto>::from_env_auto(
205                &::std::format!("{}{}_", prefix, #screaming),
206                &::std::format!("{}{}", prefix, #screaming),
207            )?
208        },
209    })
210}
211
212fn fallback_tokens(kind: &DefaultKind) -> TokenStream2 {
213    match kind {
214        DefaultKind::Standard => quote!(::std::default::Default::default()),
215        DefaultKind::Expr(expr) => quote!(#expr),
216    }
217}
218
219/// If `ty` is `Option<T>`, return `T`.
220fn option_inner(ty: &Type) -> Option<&Type> {
221    let Type::Path(tp) = ty else { return None };
222    if tp.qself.is_some() {
223        return None;
224    }
225    let seg = tp.path.segments.last()?;
226    if seg.ident != "Option" {
227        return None;
228    }
229    let PathArguments::AngleBracketed(args) = &seg.arguments else {
230        return None;
231    };
232    if args.args.len() != 1 {
233        return None;
234    }
235    match &args.args[0] {
236        GenericArgument::Type(t) => Some(t),
237        _ => None,
238    }
239}
240
241fn to_screaming_snake(s: &str) -> String {
242    let mut result = String::new();
243    let chars: Vec<char> = s.chars().collect();
244    let mut i = 0;
245    while i < chars.len() {
246        if i > 0 && chars[i].is_ascii_uppercase() {
247            let prev_lower = chars[i - 1].is_ascii_lowercase();
248            let next_lower = i + 1 < chars.len() && chars[i + 1].is_ascii_lowercase();
249            if prev_lower || (next_lower && i > 0 && chars[i - 1].is_ascii_uppercase()) {
250                result.push('_');
251            }
252        }
253        result.push(chars[i].to_ascii_uppercase());
254        i += 1;
255    }
256    result
257}
258
259struct EnvAttr {
260    rename: Option<LitStr>,
261    default: Option<DefaultKind>,
262    with: Option<Path>,
263}
264
265enum DefaultKind {
266    Standard,
267    Expr(Expr),
268}
269
270impl Parse for EnvAttr {
271    fn parse(input: ParseStream) -> syn::Result<Self> {
272        let mut rename = None;
273        let mut default = None;
274        let mut with = None;
275
276        while !input.is_empty() {
277            let ident: syn::Ident = input.parse()?;
278            match ident.to_string().as_str() {
279                "rename" => {
280                    input.parse::<Token![=]>()?;
281                    rename = Some(input.parse()?);
282                }
283                "default" => {
284                    if input.peek(Token![=]) {
285                        input.parse::<Token![=]>()?;
286                        default = Some(DefaultKind::Expr(input.parse()?));
287                    } else {
288                        default = Some(DefaultKind::Standard);
289                    }
290                }
291                "with" => {
292                    input.parse::<Token![=]>()?;
293                    let lit: LitStr = input.parse()?;
294                    let path = syn::parse_str::<Path>(&lit.value()).map_err(|e| syn::Error::new_spanned(&lit, e))?;
295                    with = Some(path);
296                }
297                other => {
298                    return Err(syn::Error::new(ident.span(), format!("unknown env attribute: `{other}`")));
299                }
300            }
301            if !input.is_empty() {
302                input.parse::<Token![,]>()?;
303            }
304        }
305
306        Ok(EnvAttr {
307            rename,
308            default,
309            with,
310        })
311    }
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317
318    #[test]
319    fn test_to_screaming_snake() {
320        assert_eq!(to_screaming_snake("url"), "URL");
321        assert_eq!(to_screaming_snake("my_field"), "MY_FIELD");
322        assert_eq!(to_screaming_snake("myField"), "MY_FIELD");
323        assert_eq!(to_screaming_snake("XMLParser"), "XML_PARSER");
324        assert_eq!(to_screaming_snake("database_url"), "DATABASE_URL");
325        assert_eq!(to_screaming_snake("db"), "DB");
326        assert_eq!(to_screaming_snake("a"), "A");
327        assert_eq!(to_screaming_snake(""), "");
328    }
329}