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 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 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 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 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 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 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 Some(EnvAttr {
138 rename: None,
139 default: Some(DefaultKind::Expr(expr)),
140 with: None,
141 }) => {
142 quote! {
143 #field_name: #expr
144 }
145 }
146 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 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 _ => {
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}