tenv/
lib.rs

1use std::{
2    collections::HashMap,
3    env, io,
4    process::{Command, ExitStatus}, path::PathBuf, ffi::OsString,
5};
6
7use clap::Parser;
8
9#[derive(Debug, Parser)]
10#[clap(author, version, about, long_about = None)]
11pub struct CommandArgs {
12    /// List of env var assignments to be set when running program
13    #[clap(short('e'), value_parser(parse_key_val::<String, String>))]
14    env_vars: Vec<(String, String)>,
15    /// Entries to add to the path variable when running program
16    #[clap(short('p'))]
17    path_additions: Vec<String>,
18    /// Program being run
19    #[clap(required(true))]
20    program: String,
21    /// Args for program to be run
22    #[clap(last(true))]
23    args: Vec<String>,
24}
25
26pub(crate) mod helpers {
27    use std::ffi::OsString;
28
29    /// Expand path if it is an argfile arg
30    pub fn expand_argfile_path(arg: String) -> OsString {
31        // If it starts with the prefix, it is an argfile arg
32        if arg.starts_with(argfile::PREFIX) {
33            // get path by itself
34            let path = arg.strip_prefix(argfile::PREFIX).expect("Already checked for this prefix");
35            // Expand path
36            if let Ok(expanded) = shellexpand::full(&path) {
37                // re-add the prefix and convert to OsString
38                return OsString::from(format!("{}{}", argfile::PREFIX, expanded))
39            }
40        }
41        // If not an argfile arg or error expanding, just return arg as an OsString
42        arg.into()
43    }
44    /// Return true if line would be a comment if in an argfile
45    pub fn is_argfile_comment(arg: &OsString) -> bool {
46        arg.to_string_lossy().starts_with('#')
47    }
48}
49
50impl CommandArgs {
51    /// Get list of args for `std::process::Command` (i.e. ["-c", `program_name`, ...`args`])
52    pub fn get_all_args() -> Result<Self, io::Error> {
53        // Expand all args and convert to OsString
54        let args_plain_iter = env::args()
55            .map(helpers::expand_argfile_path);
56        
57        // Get args and parse any from argfile if provided
58        let mut args_with_argfile = argfile::expand_args_from(
59            args_plain_iter,
60            argfile::parse_fromfile,
61            argfile::PREFIX,
62        )?;
63        
64        // Filter out comments (lines starting with '#')
65        args_with_argfile = args_with_argfile
66            .into_iter()
67            .filter(|arg| !helpers::is_argfile_comment(arg))
68            .collect();
69        
70        // Get CLI args (parse_from is from `clap::derive::Parser`)
71        let mut command_args: Self = Self::parse_from(args_with_argfile);
72        // Trim any spaces in env var name
73        command_args.env_vars.iter_mut().for_each(|(k, _v)| {
74            *k = k.trim().to_string();
75        });
76        // Trim any spaces in path
77        command_args.path_additions.iter_mut().for_each(|path| {
78            *path = path.trim().to_string();
79        });
80        Ok(command_args)
81    }
82    /// Generate new PATH by prepending path additions to existing PATH
83    fn get_prepended_path(&self, env_vars: &Option<HashMap<String, String>>) -> Option<OsString> {
84        if self.path_additions.is_empty() {
85            return None;
86        }
87        // Canonicalize paths so we can add to PATH
88        let mut path_additions: Vec<String> = self.path_additions
89            .iter()
90            .map(|s| {
91                let abso_path = dunce::simplified(&PathBuf::from(s))
92                    .to_string_lossy()
93                    // Sometimes there is white space at the beginning or end
94                    .trim()
95                    .to_string();
96                
97                // Expand ~ and other env vars
98                let expanded = shellexpand::full_with_context_no_errors::<String, _, _, PathBuf, _>(
99                    &abso_path,
100                    || {None},
101                    |var_name| {
102                        if let Some(env_map) = &env_vars {
103                            // Get from our HashMap
104                            env_map.get(var_name)
105                        }
106                        else {
107                            None
108                        }
109                    }
110                );
111                expanded.to_string()
112            })
113            .collect();
114        // get original path variable
115        let original_path: Vec<String> = env::split_paths(&env::var("PATH").unwrap_or_default())
116            // Convert paths to String
117            .map(|path| path.to_string_lossy().to_string())
118            .collect();
119
120        // Add PATH to the end of the path additions
121        path_additions.extend(original_path);
122
123        // join paths to get our new PATH environment variable
124        let new_path = env::join_paths(path_additions).expect("could not join paths");
125        Some(new_path)
126    }
127
128    fn get_env_vars(&self) -> HashMap<String, String> {
129        let mut env_vars_hashmap = HashMap::new();
130        for (var_name, value) in self.env_vars.clone() {
131            let expanded = shellexpand::full_with_context_no_errors::<String, _, _, PathBuf, _>(
132                &value, 
133                dirs::home_dir,
134                |key| {
135                    env_vars_hashmap.get(key)
136                } 
137            );
138            env_vars_hashmap.insert(var_name, expanded.to_string());
139        }
140        env_vars_hashmap
141    }
142
143    fn get_arg_list(&self) -> Vec<String> {
144        let (_, flag) = get_shell_and_flag();
145        let mut args = vec![flag.to_string(), self.program.clone()];
146        args.extend(self.args.clone());
147        args
148    }
149}
150
151/// Parse a single key-value pair
152/// taken from <https://docs.rs/clap/latest/clap/_derive/_cookbook/typed_derive/index.html>
153fn parse_key_val<T, U>(
154    s: &str,
155) -> Result<(T, U), Box<dyn std::error::Error + Send + Sync + 'static>>
156where
157    T: std::str::FromStr,
158    T::Err: std::error::Error + Send + Sync + 'static,
159    U: std::str::FromStr,
160    U::Err: std::error::Error + Send + Sync + 'static,
161{
162    let pos = s
163        .find('=')
164        .ok_or_else(|| format!("invalid KEY=value: no `=` found in `{}`", s))?;
165    Ok((s[..pos].parse()?, s[pos + 1..].parse()?))
166}
167
168/// ctrl+c handler so that tenv itself can't be interrupted.
169/// Commands run by tenv can still be cancelled, therefore ending the execution of tenv.
170/// Not doing this causes problems with starship on powershell.
171pub fn init_ctrlc_handler() -> Result<(), ctrlc::Error> {
172    // Set empty ctrl+c handler, just so doesn't stop tenv itself when ctrl+c entered
173    ctrlc::set_handler(move || {
174        // println!("ctrl+c pressed");
175    })
176}
177
178/// Runs command, while setting environment variables before unsets them after command is completed
179pub fn run(command_args: &CommandArgs) -> Result<ExitStatus, io::Error> {
180    // Convert Vec of env vars to HashMap
181    let hash_map_vars: HashMap<String, String> = command_args.get_env_vars();
182
183    // Get shell and appropriate flag to run command through OS shell
184    let (shell, _) = get_shell_and_flag();
185    // Combine flag with program and its args [flag, program name, rest of args]
186    let final_args = command_args.get_arg_list();
187
188    // Build command with shell
189    let mut command = Command::new(shell);
190    // Add args for shell to run command
191    command
192        .args(&final_args)
193        // Set env variables
194        .envs(&hash_map_vars);
195    
196    // If path_additions passed to CLI, get and set to new path
197    if let Some(new_path) = command_args.get_prepended_path(&Some(hash_map_vars)) {
198        // Set PATH env var
199        command.env("PATH", new_path);
200    }
201
202    // Run command and return status
203    command.status()
204}
205
206
207
208/// Return ("powershell", "-Command") for windows or ("bash", "-c") for any other OS
209const fn get_shell_and_flag<'a>() -> (&'a str, &'a str) {
210    if cfg!(windows) {
211        ("powershell", "-Command")
212    } else {
213        ("bash", "-c")
214    }
215}
tenv/lib.rs

tenv/
lib.rs
