const_env_impl__value/
lib.rs

1#![allow(non_snake_case)]
2
3use std::collections::HashMap;
4
5use proc_macro2::TokenStream;
6use quote::quote_spanned;
7use syn::{Expr, ExprLit, Lit};
8use syn::spanned::Spanned;
9
10pub trait ReadEnv {
11    fn read_env(&self, var_name: &String) -> Option<String>;
12}
13
14pub struct RealEnv;
15
16impl ReadEnv for RealEnv {
17    fn read_env(&self, var_name: &String) -> Option<String> {
18        std::env::var_os(var_name)
19            .and_then(|s| s.into_string().ok())
20    }
21}
22
23pub struct TestEnv {
24    env_vars: HashMap<String, String>
25}
26
27impl TestEnv {
28    pub fn builder() -> TestEnvBuilder {
29        TestEnvBuilder {
30            env_vars: HashMap::new()
31        }
32    }
33}
34
35impl ReadEnv for TestEnv {
36    fn read_env(&self, var_name: &String) -> Option<String> {
37        self.env_vars.get(var_name).cloned()
38    }
39}
40
41pub struct TestEnvBuilder {
42    env_vars: HashMap<String, String>
43}
44
45impl TestEnvBuilder {
46    pub fn set(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
47        self.env_vars.insert(name.into(), value.into());
48        self
49    }
50
51    pub fn build(self) -> TestEnv {
52        TestEnv {
53            env_vars: self.env_vars
54        }
55    }
56}
57
58/// Inner implementation details of `const_env::from_env`.
59pub fn from_env(attr: TokenStream, item: TokenStream, read_env: impl ReadEnv) -> TokenStream {
60    if let Ok(mut item_const) = syn::parse2::<syn::ItemConst>(item.clone()) {
61        let default_var_name = format!("{}", item_const.ident);
62        let var_name = extract_var_name(attr, default_var_name);
63        let var_value = match read_env.read_env(&var_name) {
64            Some(val) => val,
65            None => return item
66        };
67        let new_expr = value_to_literal(&var_value, &item_const.expr);
68        let span = item_const.span();
69        item_const.expr = Box::new(new_expr);
70        quote_spanned!(span => #item_const)
71    } else if let Ok(mut item_static) = syn::parse2::<syn::ItemStatic>(item.clone()) {
72        let default_var_name = format!("{}", item_static.ident);
73        let var_name = extract_var_name(attr, default_var_name);
74        let var_value = match read_env.read_env(&var_name) {
75            Some(val) => val,
76            None => return item
77        };
78        let new_expr = value_to_literal(&var_value, &item_static.expr);
79        let span = item_static.span();
80        item_static.expr = Box::new(new_expr);
81        quote_spanned!(span => #item_static)
82    } else {
83        panic!("TODO: error reporting");
84    }
85}
86
87/// Inner implementation details of `const_env::value_from_env`.
88pub fn value_from_env(item: TokenStream, read_env: impl ReadEnv) -> TokenStream {
89    let expr_type: syn::ExprType = syn::parse2(item)
90        .expect("Unable to parse environment variable name as expression");
91    let var_name = match &*expr_type.expr {
92        Expr::Lit(syn::ExprLit { lit: syn::Lit::Str(var_name), .. }) => var_name.value(),
93        invalid_expr => panic!("expected a string literal containing an environment variable name but got '{:?}'", invalid_expr),
94    };
95
96    let var_value = match read_env.read_env(&var_name) {
97        Some(val) => val,
98        None => panic!("no value available for environment variable '{}'", var_name),
99    };
100
101    let span = expr_type.span();
102
103    let type_name = match &*expr_type.ty {
104        syn::Type::Path(p) => {
105            match p.path.get_ident() {
106                Some(ident) => ident.to_string(),
107                None => panic!("expected a type with no path separators but got '{:?}'", p),
108            }
109        },
110        unsupported_ty => panic!("unsupported type for environment variable substitution: {:?}", unsupported_ty),
111    };
112
113    let new_expr = match &type_name[..] {
114        "u8" | "u16" | "u32" | "u64" | "usize" |
115            "i8" | "i16" | "i32" | "i64" | "isize" => {
116            syn::parse2::<syn::LitInt>(format!("{}{}", var_value, type_name).parse().unwrap()).unwrap() // add type suffix on interger literals
117        },
118        unsupported_type => panic!("unsupported type for environment variable substitution: '{}'", unsupported_type),
119    };
120    quote_spanned!(span => #new_expr)
121}
122
123fn extract_var_name(attr: TokenStream, default: String) -> String {
124    if attr.is_empty() {
125        return default;
126    }
127    let expr: Expr = syn::parse2(attr)
128        .expect("Unable to parse attribute args as expression");
129    extract_var_name_from_expr(&expr)
130}
131
132fn extract_var_name_from_expr(expr: &Expr) -> String {
133    match expr {
134        Expr::Lit(literal) => {
135            match &literal.lit {
136                Lit::Str(lit_str) => {
137                    lit_str.value()
138                },
139                _ => panic!("Attribute arguments are not a valid string literal")
140            }
141        },
142        Expr::Paren(paren) => {
143            extract_var_name_from_expr(&paren.expr)
144        },
145        _ => {
146            panic!("Attribute arguments are not a valid string literal expression: {:?}", expr)
147        }
148    }
149}
150
151fn value_to_literal(value: &str, original_expr: &Expr) -> Expr {
152    match original_expr {
153        Expr::Unary(_) => {
154            // A unary sign indicates this is a numeric literal which doesn't need any
155            // escaping, so we can parse it directly.
156            let new: Expr = syn::parse_str(value)
157                .expect("Failed to parse environment variable contents as valid expresion");
158            return new;
159        },
160        Expr::Lit(literal) => {
161            let new_lit = match &literal.lit {
162                Lit::Str(original) => {
163                    let mut new: syn::LitStr = syn::parse_str(&format!("\"{}\"", value))
164                        .expect("Failed to parse environment variable contents as literal string");
165                    new.set_span(original.span());
166                    Lit::Str(new)
167                },
168                Lit::ByteStr(original) => {
169                    let mut new: syn::LitByteStr = syn::parse_str(&format!("b\"{}\"", value))
170                        .expect("Failed to parse environment variable contents as literal byte string");
171                    new.set_span(original.span());
172                    Lit::ByteStr(new)
173                },
174                Lit::Byte(original) => {
175                    let mut new: syn::LitByte = syn::parse_str(&format!("b'{}'", value))
176                        .expect("Failed to parse environment variable contents as literal byte");
177                    new.set_span(original.span());
178                    Lit::Byte(new)
179                },
180                Lit::Char(original) => {
181                    let mut new: syn::LitChar = syn::parse_str(&format!("'{}'", value))
182                        .expect("Failed to parse environment variable contents as literal character");
183                    new.set_span(original.span());
184                    Lit::Char(new)
185                },
186                // These variants do not need any escaping and can be parsed as an expression
187                // directly.
188                Lit::Bool(_) | Lit::Int(_) | Lit::Float(_) | Lit::Verbatim(_) => {
189                    let new: Expr = syn::parse_str(value)
190                        .expect("Failed to parse environment variable contents as valid expression");
191                    return new;
192                }
193            };
194            ExprLit {
195                attrs: literal.attrs.clone(),
196                lit: new_lit
197            }.into()
198        },
199        _ => panic!("Original const expression was not a recognized literal expression")
200    }
201}
const_env_impl__value/lib.rs

const_env_impl__value/
lib.rs
