diff options
Diffstat (limited to 'zkvms_host_io/src/lib.rs')
| -rw-r--r-- | zkvms_host_io/src/lib.rs | 19 |
1 files changed, 18 insertions, 1 deletions
diff --git a/zkvms_host_io/src/lib.rs b/zkvms_host_io/src/lib.rs index a3bdb1b..c029975 100644 --- a/zkvms_host_io/src/lib.rs +++ b/zkvms_host_io/src/lib.rs @@ -2,7 +2,7 @@ use clap::{Parser, ValueEnum}; use num_traits::NumCast; use serde::{ Serialize, Deserialize }; use env_file_reader::read_str; -use std::{ fs::read_to_string, collections::HashMap }; +use std::{env, option::Option, fs::read_to_string, collections::HashMap}; pub use input_macros::foreach_input_field; static DEFAULT_PUBLIC_INPUT: &str = include_str!(concat!(env!("INPUTS_DIR"), "/default_public_input.toml")); @@ -31,6 +31,23 @@ pub struct RunWith<T> { pub default_env: HashMap<String, String>, } +impl<T> RunWith<T> { + pub fn env_then_or<U>(&self, variable_name: &str, then_apply: fn(String) -> Option<U>, else_const: U) -> U { + env::var(variable_name) + .ok() + .and_then(then_apply) + .unwrap_or(self + .default_env + .get(variable_name) + .and_then(|x| then_apply(x.clone())) + .unwrap_or(else_const)) + } + + pub fn env_or(&self, variable_name: &str, else_const: &str) -> String { + self.env_then_or(variable_name, |x| Some(x), else_const.to_string()) + } +} + input_macros::generate_input_struct!(); pub fn read_args() -> RunWith<Input> { |