1use std::borrow::Borrow;
22use std::collections::HashMap;
23use std::hash::{BuildHasher, Hash};
24
25use proc_macro::TokenStream;
26use proc_macro2::Span;
27use proc_macro2::TokenStream as TokenStream2;
28use quote::quote_spanned;
29use syn::parse::{Parse, ParseStream};
30use syn::punctuated::Punctuated;
31use syn::spanned::Spanned;
32use syn::{parse_macro_input, Expr, Ident, Token, Type};
33
34struct EnvironmentVariable {
35 name: Ident,
36 ty: Type,
37 default: Option<Expr>,
38}
39
40impl Parse for EnvironmentVariable {
41 fn parse(input: ParseStream) -> syn::Result<Self> {
42 let name: Ident = input.parse()?;
43 input.parse::<Token![:]>()?;
44 let ty: Type = input.parse()?;
45 let default = if input.peek(Token![=]) {
46 input.parse::<Token![=]>()?;
47 Some(input.parse()?)
48 } else {
49 None
50 };
51
52 Ok(Self { name, ty, default })
53 }
54}
55
56struct EnvironmentVariables(Punctuated<EnvironmentVariable, Token![,]>);
57
58impl Parse for EnvironmentVariables {
59 fn parse(input: ParseStream) -> syn::Result<Self> {
60 Ok(Self(Punctuated::parse_terminated(input)?))
61 }
62}
63
64impl IntoIterator for EnvironmentVariables {
65 type Item = EnvironmentVariable;
66 type IntoIter = <Punctuated<EnvironmentVariable, Token![,]> as IntoIterator>::IntoIter;
67
68 fn into_iter(self) -> Self::IntoIter {
69 self.0.into_iter()
70 }
71}
72
73trait HashMapExt<K, V, S> {
74 fn get_any_case(&self, key: &str) -> Option<&V>
75 where
76 K: Borrow<str>;
77}
78
79impl<K, V, S> HashMapExt<K, V, S> for HashMap<K, V, S>
80where
81 K: Hash + Eq,
82 S: BuildHasher,
83{
84 fn get_any_case(&self, key: &str) -> Option<&V>
85 where
86 K: Borrow<str>,
87 {
88 match self.get(key) {
89 Some(value) => Some(value),
90 None => {
91 let key_lower = key.to_lowercase();
92 for key in self.keys() {
93 if key.borrow().to_lowercase() == key_lower {
94 return self.get(key.borrow());
95 }
96 }
97
98 None
99 }
100 }
101 }
102}
103
104fn load_env() -> Result<HashMap<String, String>, dotenvy::Error> {
105 let mut result = HashMap::new();
106
107 for item in dotenvy::dotenv_iter()? {
108 let (key, value) = item?;
109 result.insert(key, value);
110 }
111
112 Ok(result)
113}
114
115fn is_str_literal(ty: &syn::Type) -> bool {
116 if let syn::Type::Reference(syn::TypeReference { elem, .. }) = ty {
117 if let syn::Type::Path(syn::TypePath { path, .. }) = elem.as_ref() {
118 if path.segments.len() == 1 && path.segments[0].ident == "str" {
119 return true;
120 }
121 }
122 }
123
124 false
125}
126
127fn expand(variables: impl IntoIterator<Item = EnvironmentVariable>) -> syn::Result<TokenStream2> {
128 let env_map = load_env().map_err(|err| {
129 syn::Error::new(
130 Span::call_site(),
131 format!("Failed to load environment: {}", err),
132 )
133 })?;
134
135 let mut result = Vec::new();
136 for EnvironmentVariable { name, ty, default } in variables {
137 let value = match (env_map.get_any_case(&name.to_string()), default) {
138 (Some(value), _) => {
139 if is_str_literal(&ty) {
140 syn::parse_str(&format!("\"{}\"", value))?
141 } else {
142 syn::parse_str::<syn::Expr>(value)?
143 }
144 }
145 (None, Some(value)) => value,
146 (None, None) => {
147 return Err(syn::Error::new(
148 name.span(),
149 format!("Environment variable '{}' not found", name),
150 ));
151 }
152 };
153
154 result.push(quote_spanned! {
155 ty.span() => { let _res: #ty = #value; _res }
156 });
157 }
158
159 Ok(quote::quote! {
160 (#(#result),*)
161 })
162}
163
164#[proc_macro]
170pub fn dotenvy(input: TokenStream) -> TokenStream {
171 expand(parse_macro_input!(input as EnvironmentVariables))
172 .unwrap_or_else(syn::Error::into_compile_error)
173 .into()
174}