diff --git a/crates/rattler_shell/src/activation.rs b/crates/rattler_shell/src/activation.rs index 1680015510..5bc0044648 100644 --- a/crates/rattler_shell/src/activation.rs +++ b/crates/rattler_shell/src/activation.rs @@ -16,6 +16,7 @@ use std::{ use anyhow::{Context, Result}; use fs_err as fs; use indexmap::IndexMap; +use itertools::Itertools; use rattler_conda_types::Platform; #[cfg(target_family = "unix")] use rattler_pty::unix::PtySession; @@ -93,9 +94,14 @@ pub struct Activator { /// A list of scripts to run when deactivating the environment pub deactivation_scripts: Vec, - /// A list of environment variables to set when activating the environment + /// A list of environment variables to set before running the activation + /// scripts. These are evaluated before `activation_scripts` have run. pub env_vars: IndexMap, + /// A list of environment variables to set after running the activation + /// scripts. These are evaluated after `activation_scripts` have run. + pub post_activation_env_vars: IndexMap, + /// The platform for which to generate the Activator pub platform: Platform, } @@ -323,6 +329,17 @@ pub struct ActivationResult { } impl Activator { + /// Return unique env var keys from both `env_vars` and `post_activation_env_vars` in insertion order. + fn unique_env_keys(&self) -> impl Iterator { + self.env_vars + .keys() + .chain(self.post_activation_env_vars.keys()) + .map(String::as_str) + .unique() + } + + // moved: apply_env_vars_with_backup now lives on `ShellScript` + /// Create a new activator for the given conda environment. /// /// # Arguments @@ -368,6 +385,7 @@ impl Activator { activation_scripts, deactivation_scripts, env_vars, + post_activation_env_vars: IndexMap::new(), platform, }) } @@ -496,22 +514,19 @@ impl Activator { script.set_env_var("CONDA_PREFIX", &self.target_prefix.to_string_lossy())?; // For each environment variable that was set during activation - for (key, value) in &self.env_vars { - // Save original value if it exists - if let Some(existing_value) = variables.current_env.get(key) { - script.set_env_var( - &format!("CONDA_ENV_SHLVL_{new_shlvl}_{key}"), - existing_value, - )?; - } - // Set new value - script.set_env_var(key, value)?; - } + script.apply_env_vars_with_backup(&variables.current_env, new_shlvl, &self.env_vars)?; for activation_script in &self.activation_scripts { script.run_script(activation_script)?; } + // Set environment variables that should be applied after activation scripts + script.apply_env_vars_with_backup( + &variables.current_env, + new_shlvl, + &self.post_activation_env_vars, + )?; + Ok(ActivationResult { script, path }) } @@ -539,8 +554,8 @@ impl Activator { "Proceeding to unset conda variables without restoring previous values.", )?; - // Just unset without restoring - for (key, _) in &self.env_vars { + // Just unset without restoring (each key once) + for key in self.unique_env_keys() { script.unset_env_var(key)?; } script.unset_env_var("CONDA_PREFIX")?; @@ -553,8 +568,8 @@ impl Activator { "Proceeding to unset conda variables without restoring previous values.", )?; - // Just unset without restoring - for (key, _) in &self.env_vars { + // Just unset without restoring (each key once) + for key in self.unique_env_keys() { script.unset_env_var(key)?; } script.unset_env_var("CONDA_PREFIX")?; @@ -563,7 +578,7 @@ impl Activator { Some(current_level) => { // Unset the current level // For each environment variable that was set during activation - for (key, _) in &self.env_vars { + for key in self.unique_env_keys() { let backup_key = format!("CONDA_ENV_SHLVL_{current_level}_{key}"); script.restore_env_var(key, &backup_key)?; } @@ -613,9 +628,8 @@ impl Activator { ShellScript::new(self.shell_type.clone(), self.platform); activation_detection_script .print_env()? - .echo(ENV_START_SEPARATOR)?; - activation_detection_script.append_script(&activation_script); - activation_detection_script + .echo(ENV_START_SEPARATOR)? + .append_script(&activation_script) .echo(ENV_START_SEPARATOR)? .print_env()?; @@ -686,6 +700,76 @@ mod tests { use crate::activation::PathModificationBehavior; use crate::shell::{self, native_path_to_unix, ShellEnum}; + #[test] + #[cfg(unix)] + fn test_post_activation_env_vars_applied_after_scripts_bash() { + let temp_dir = TempDir::new("test_post_activation_env_vars").unwrap(); + + // Create a dummy activation script so the activator will run it + let activate_dir = temp_dir.path().join("etc/conda/activate.d"); + fs::create_dir_all(&activate_dir).unwrap(); + let script_path = activate_dir.join("script1.sh"); + fs::write(&script_path, "# noop\n").unwrap(); + + // Build an activator with both pre and post env vars + let pre_env = IndexMap::from_iter([(String::from("A"), String::from("x"))]); + + // Ensure we also override a pre var in post + let post_env = IndexMap::from_iter([ + (String::from("B"), String::from("y")), + (String::from("A"), String::from("z")), + ]); + + let activator = Activator { + target_prefix: temp_dir.path().to_path_buf(), + shell_type: shell::Bash, + paths: vec![temp_dir.path().join("bin")], + activation_scripts: vec![script_path.clone()], + deactivation_scripts: vec![], + env_vars: pre_env, + post_activation_env_vars: post_env, + platform: Platform::current(), + }; + + let result = activator + .activation(ActivationVariables { + conda_prefix: None, + path: None, + path_modification_behavior: PathModificationBehavior::Prepend, + current_env: HashMap::new(), + }) + .unwrap(); + + let mut contents = result.script.contents().unwrap(); + + // Normalize prefix path for consistent assertions + let prefix = temp_dir.path().to_str().unwrap(); + contents = contents.replace(prefix, "__PREFIX__"); + + // Check ordering: pre env vars before script run, post env vars after script run + let idx_pre_a = contents.find("export A=x").expect("missing pre env A=x"); + let idx_run = contents + .find(". __PREFIX__/etc/conda/activate.d/script1.sh") + .expect("missing activation script run"); + let idx_post_b = contents.find("export B=y").expect("missing post env B=y"); + let idx_post_a = contents + .find("export A=z") + .expect("missing post override A=z"); + + assert!( + idx_pre_a < idx_run, + "pre env var should be before activation script" + ); + assert!( + idx_run < idx_post_b, + "post env var should be after activation script" + ); + assert!( + idx_run < idx_post_a, + "post override should be after activation script" + ); + } + #[test] fn test_collect_scripts() { let tdir = TempDir::new("test").unwrap(); @@ -1037,6 +1121,7 @@ mod tests { activation_scripts: vec![], deactivation_scripts: vec![], env_vars: env_vars.clone(), + post_activation_env_vars: IndexMap::new(), platform: Platform::current(), }; @@ -1093,6 +1178,7 @@ mod tests { activation_scripts: vec![], deactivation_scripts: vec![], env_vars: env_vars.clone(), + post_activation_env_vars: IndexMap::new(), platform: Platform::current(), }; @@ -1162,6 +1248,7 @@ mod tests { activation_scripts: vec![], deactivation_scripts: vec![], env_vars: second_env_vars.clone(), + post_activation_env_vars: IndexMap::new(), platform: Platform::current(), }; @@ -1280,6 +1367,7 @@ mod tests { activation_scripts: vec![], deactivation_scripts: vec![], env_vars: second_env_vars.clone(), + post_activation_env_vars: IndexMap::new(), platform: Platform::current(), }; diff --git a/crates/rattler_shell/src/shell/mod.rs b/crates/rattler_shell/src/shell/mod.rs index 7a652fe995..32244bd4dc 100644 --- a/crates/rattler_shell/src/shell/mod.rs +++ b/crates/rattler_shell/src/shell/mod.rs @@ -12,6 +12,7 @@ use std::{ }; use enum_dispatch::enum_dispatch; +use indexmap::IndexMap; use itertools::Itertools; use rattler_conda_types::Platform; use thiserror::Error; @@ -1046,6 +1047,26 @@ impl ShellScript { } } + /// Apply the provided environment variables to the script while + /// backing up existing values to the current shell level. + pub fn apply_env_vars_with_backup( + &mut self, + current_env: &HashMap, + new_shlvl: i32, + envs: &IndexMap, + ) -> Result<&mut Self, ShellError> { + for (key, value) in envs { + if let Some(existing_value) = current_env.get(key) { + self.set_env_var( + &format!("CONDA_ENV_SHLVL_{new_shlvl}_{key}"), + existing_value, + )?; + } + self.set_env_var(key, value)?; + } + Ok(self) + } + /// Export an environment variable. pub fn set_env_var(&mut self, env_var: &str, value: &str) -> Result<&mut Self, ShellError> { self.shell.set_env_var(&mut self.contents, env_var, value)?;