1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

extern crate proc_macro;

use quote::quote;
use syn::{
    parse::{Parse, ParseStream},
    parse_macro_input, LitStr, Token,
};

struct DefaultEnv {
    env_var: LitStr,
    default_value: proc_macro2::TokenStream,
}

impl Parse for DefaultEnv {
    fn parse(input: ParseStream) -> syn::parse::Result<Self> {
        let env_var = input.parse()?;
        input.parse::<Token![,]>()?;
        let default_value = input.parse()?;
        Ok(Self {
            env_var,
            default_value,
        })
    }
}

#[proc_macro]
pub fn default_env(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let DefaultEnv {
        env_var,
        default_value,
    } = parse_macro_input!(input as DefaultEnv);

    let var_or_default = match std::env::var(env_var.value()) {
        Ok(var) => quote! { #var },
        Err(_) => quote! { #default_value },
    };

    proc_macro::TokenStream::from(var_or_default)
}