Skip to main content

lexe_common/
or_env.rs

1//! Use `OrEnvExt` when you have a cli arg (e.g. `db_url: Option<String>`) that
2//! can also be set by a "fallback" env variable (e.g. `$DATABASE_URL`). Then,
3//! when initializing the args, just use `db_url.or_env_mut("DATABASE_URL")?`.
4//!
5//! This works with any `FromStr` type, in which case it will parse the env
6//! value and return an error if that fails.
7
8use std::{env, str::FromStr};
9
10use anyhow::Context;
11
12pub trait OrEnvExt: Sized {
13    /// Analogous to [`Option::or_else`]. Takes ownership of the arg if set,
14    /// otherwise initializes the arg from env. Used to initialize args lazily.
15    /// Also has good ergonomics with [`Result`] / [`Option`] chains.
16    fn or_env(mut self, env_var: &'static str) -> anyhow::Result<Self> {
17        self.or_env_mut(env_var)?;
18        Ok(self)
19    }
20
21    /// If the arg is not set, initialize the arg from env by mutating the arg
22    /// in place. Used to proactively initialize and validate args.
23    fn or_env_mut(
24        &mut self,
25        env_var: &'static str,
26    ) -> anyhow::Result<&mut Self>;
27}
28
29fn env_var_opt(env_var: &'static str) -> anyhow::Result<Option<String>> {
30    match env::var(env_var) {
31        Ok(val_str) => Ok(Some(val_str)),
32        Err(env::VarError::NotPresent) => Ok(None),
33        Err(env::VarError::NotUnicode(s)) =>
34            Err(anyhow::format_err!("invalid unicode: '{:?}'", s)),
35    }
36}
37
38impl<T> OrEnvExt for Option<T>
39where
40    T: FromStr,
41    T::Err: Into<anyhow::Error>,
42{
43    fn or_env_mut(
44        &mut self,
45        env_var: &'static str,
46    ) -> anyhow::Result<&mut Option<T>> {
47        if self.is_none() {
48            // If no env var, do nothing. Error if not UTF-8 encoded.
49            let val_str = match env_var_opt(env_var).context(env_var)? {
50                Some(v) => v,
51                None => return Ok(self),
52            };
53            let val = T::from_str(&val_str)
54                .map_err(Into::into)
55                .with_context(|| format!("Invalid env value `${env_var}`"))?;
56            *self = Some(val);
57        }
58
59        Ok(self)
60    }
61}
62
63impl OrEnvExt for bool {
64    fn or_env_mut(
65        &mut self,
66        env_var: &'static str,
67    ) -> anyhow::Result<&mut bool> {
68        if !*self {
69            // If no env var, do nothing. Error if not UTF-8 encoded.
70            let val_str = match env_var_opt(env_var).context(env_var)? {
71                Some(v) => v,
72                None => return Ok(self),
73            };
74            let val = bool::from_str(&val_str)
75                .with_context(|| format!("Invalid env value `${env_var}`"))?;
76            *self = val;
77        }
78
79        Ok(self)
80    }
81}