1mod attr;
2
3use crate::attr::{EnvAttr, EnvAttrOp};
4use proc_macro::TokenStream;
5use quote::{ToTokens, format_ident, quote, quote_spanned};
6use syn::spanned::Spanned;
7use syn::{Data, DeriveInput, Fields, Type, parse_macro_input, parse_quote};
8
9#[proc_macro_attribute]
10pub fn serde_with_env(_: TokenStream, item: TokenStream) -> TokenStream {
11 let mut derive_input: DeriveInput = parse_macro_input!(item);
12
13 let struct_name = &derive_input.ident;
14
15 let mod_name = format_ident!("__serde_with_env__{}", struct_name);
16 let shadow_struct_name = format_ident!("__{}", struct_name);
17 let try_from_path = format!("{mod_name}::{shadow_struct_name}");
18
19 derive_input
20 .attrs
21 .push(parse_quote!( #[serde(try_from = #try_from_path)]));
22
23 quote_spanned! { derive_input.span() =>
24 #[derive(serde_with_env::SerdeWithEnv)]
25 #derive_input
26 }
27 .into_token_stream()
28 .into()
29}
30
31#[proc_macro_derive(SerdeWithEnv, attributes(with_env))]
36pub fn serde_with_env_derive(item: TokenStream) -> TokenStream {
37 let mut derive_input: DeriveInput = parse_macro_input!(item);
38 let struct_name = &derive_input.ident;
39
40 let mod_name = format_ident!("__serde_with_env__{}", struct_name);
41 let shadow_struct_name = format_ident!("__{}", struct_name);
42
43 let mut generated_get_env = Vec::new();
44 let mut generated_from_fields = Vec::new();
45 let mut generated_fields = Vec::new();
46
47 if let Data::Struct(ref mut data) = derive_input.data {
48 if let Fields::Named(ref mut fields) = data.fields {
49 for field in fields.named.iter_mut() {
50 let mut field_is_option = false;
51 let field_name = field.ident.as_ref().unwrap();
52 let maybe_attr = match attr::get_env_attr(&field.attrs) {
53 None => {
54 generated_fields.push(field.to_token_stream());
55 generated_from_fields.push(quote! {
56 #field_name: v.#field_name
57 });
58 continue;
59 }
60 Some(maybe_attr) => {
61 field.attrs.retain(|attr| !attr.path().is_ident("with_env"));
62
63 let mut field = field.clone();
64 let ty = field.ty.clone();
65
66 if let Type::Path(path) = &ty
67 && let Some(segment) = path.path.segments.first()
68 && segment.ident == "Option"
69 {
70 field_is_option = true;
71 }
72
73 if !field_is_option {
74 field.ty = parse_quote!(Option<#ty>);
75 }
76
77 if let Ok(EnvAttr {
78 op: EnvAttrOp::Or | EnvAttrOp::Over,
79 ..
80 }) = &maybe_attr
81 {
82 generated_fields.push(field.to_token_stream());
83 }
84 generated_from_fields.push(quote! {
85 #field_name
86 });
87 maybe_attr
88 }
89 };
90
91 match maybe_attr {
92 Ok(env_attr) => {
93 let env_name = env_attr.name.value();
94 let env_default = env_attr.default.as_ref();
95 let env_convert = env_attr.convert.as_ref();
96
97 let missing_err = format!("Missing \"{env_name}\" environment variable.");
98 let parse_err =
99 format!("Cannot parse \"{env_name}\" environment variable: {{err}}");
100
101 let field_as_some = match field_is_option {
102 true => quote! { Ok(Some(v)) },
103 false => quote! { Ok(v) },
104 };
105
106 let parse_as_some = match (field_is_option, env_convert) {
107 (true, None) => {
108 quote! { v.parse().map(Some).map_err(|err| format!(#parse_err)) }
109 }
110 (false, None) => {
111 quote! { v.parse().map_err(|err| format!(#parse_err)) }
112 }
113 (_, Some(convert)) => {
114 let ident = format_ident!("{}", convert.value());
115 quote! { #ident(v) }
116 }
117 };
118
119 let not_present_error = match (field_is_option, env_default) {
120 (true, Some(default)) => quote! { Ok(Some(#default.into())) },
121 (false, Some(default)) => quote! { Ok(#default.into()) },
122 (true, None) => quote! { Ok(None)},
123 (false, None) => quote! { Err(#missing_err.to_string()) },
124 };
125
126 let env_strategy = match env_attr.op {
127 EnvAttrOp::Or => {
128 quote! {
129 let #field_name = if let Some(v) = v.#field_name {
130 #field_as_some
131 } else {
132 match std::env::var(#env_name) {
133 Ok(v) => #parse_as_some,
134 Err(err) => #not_present_error,
135 }
136 }?;
137 }
138 }
139 EnvAttrOp::Only => {
140 quote! {
141 let #field_name = match std::env::var(#env_name) {
142 Ok(v) => #parse_as_some,
143 Err(err) => #not_present_error,
144 }?;
145 }
146 }
147 EnvAttrOp::Over => {
148 let not_present_error = match (field_is_option, env_default) {
149 (false, None) => quote! { Err(#missing_err.to_string()) },
150 _ => not_present_error,
151 };
152
153 quote! {
154 let #field_name = match std::env::var(#env_name) {
155 Ok(v) => #parse_as_some,
156 Err(_) => match v.#field_name {
157 None => #not_present_error,
158 Some(v) => #field_as_some,
159 },
160 }?;
161 }
162 }
163 };
164
165 generated_get_env.push(env_strategy);
166 }
167 Err(err) => {
168 let err = err.into_compile_error();
169 generated_get_env.push(quote! {
170 compile_error!(#err);
171 });
172 }
173 }
174 }
175 }
176 }
177
178 let try_from_impl = quote! {
179 #[automatically_derived]
180 impl TryFrom<#shadow_struct_name> for #struct_name {
181 type Error = String;
182 fn try_from(v: #shadow_struct_name) -> Result<Self, Self::Error> {
183 #(#generated_get_env)*
184
185 Ok(Self {
186 #(#generated_from_fields),*
187 })
188 }
189 }
190 };
191
192 quote! {
193 #[allow(non_snake_case)]
194 mod #mod_name {
195 use super::*;
196
197 #[derive(serde::Deserialize)]
198 pub struct #shadow_struct_name {
199 #(#generated_fields),*
200 }
201
202 #try_from_impl
203 }
204 }
205 .into()
206}