Skip to main content

compile_dotenv/
lib.rs

1//! Code for the [`compile_env!`] macro.
2use dotenv::dotenv;
3use proc_macro::{TokenStream, TokenTree};
4use quote::quote;
5use std::env;
6
7/// Returns an environment variable in the `.env` file or in the environment.
8///
9/// # Panics
10/// Panics if there are no arguments or more than two arguments.
11#[proc_macro]
12pub fn compile_env(input: TokenStream) -> TokenStream {
13    let tokens: Vec<_> = input.into_iter().collect();
14
15    let (name, default): (String, String) = match tokens.as_slice() {
16        [TokenTree::Literal(lit)] => (unwrap_string_literal(lit), String::new()),
17        [
18            TokenTree::Literal(lit),
19            TokenTree::Punct(punct),
20            TokenTree::Literal(default_lit),
21        ] if punct.as_char() == ',' => {
22            (unwrap_string_literal(lit), unwrap_string_literal(default_lit))
23        },
24        _ => panic!("This macro only accepts one or two arguments"),
25    };
26
27    // Read the .env file
28    let _ = dotenv();
29
30    let value = env::var(name).unwrap_or_else(|_| default.to_string());
31    quote!(#value).into()
32}
33
34/// Gets a [`String`] from a [`proc_macro::Literal`].
35///
36/// Inspired from <https://docs.rs/include_dir_macros/0.7.4/src/include_dir_macros/lib.rs.html#31>.
37///
38/// # Panics
39/// Panics if the [`proc_macro::Literal`] doesn't contain a single string argument.
40fn unwrap_string_literal(lit: &proc_macro::Literal) -> String {
41    let mut repr = lit.to_string();
42    assert!(
43        repr.starts_with('"') && repr.ends_with('"'),
44        "This macro only accepts a single, non-empty string argument"
45    );
46    repr.remove(0);
47    repr.pop();
48    repr
49}