From 1dc0beb8c8f5c89f0fc7153540f50bfb3cf96760 Mon Sep 17 00:00:00 2001 From: Bas Zalmstra Date: Fri, 30 Jun 2023 17:00:45 +0200 Subject: [PATCH] feat: run activation and capture result (#239) --- crates/rattler_shell/Cargo.toml | 6 +- crates/rattler_shell/src/activation.rs | 188 +++++++++++++++++- crates/rattler_shell/src/shell/mod.rs | 42 +++- ...attler_shell__activation__tests__bash.snap | 9 + ...ler_shell__activation__tests__cmd.exe.snap | 9 + ...attler_shell__activation__tests__pwsh.snap | 9 + ...rattler_shell__activation__tests__zsh.snap | 9 + 7 files changed, 264 insertions(+), 8 deletions(-) create mode 100644 crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__bash.snap create mode 100644 crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__cmd.exe.snap create mode 100644 crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__pwsh.snap create mode 100644 crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__zsh.snap diff --git a/crates/rattler_shell/Cargo.toml b/crates/rattler_shell/Cargo.toml index 53174e2b8..2e2566a76 100644 --- a/crates/rattler_shell/Cargo.toml +++ b/crates/rattler_shell/Cargo.toml @@ -16,10 +16,12 @@ indexmap = "1.9.3" itertools = "0.10.5" rattler_conda_types = { version = "0.5.0", path = "../rattler_conda_types" } serde_json = { version = "1.0.96", features = ["preserve_order"] } +shlex = "1.1.0" +sysinfo = { version = "0.29.2", optional = true } +tempfile = "3.5.0" thiserror = "1.0.40" tracing = "0.1.37" -sysinfo = { version = "0.29.2", optional = true } [dev-dependencies] +insta = { version = "1.29.0", features = ["yaml"] } tempdir = "0.3.7" -insta = "1.29.0" diff --git a/crates/rattler_shell/src/activation.rs b/crates/rattler_shell/src/activation.rs index 3ffa88510..535ac3c50 100644 --- a/crates/rattler_shell/src/activation.rs +++ b/crates/rattler_shell/src/activation.rs @@ -2,7 +2,9 @@ //! This crate provides helper functions to activate and deactivate virtual environments. +use std::collections::HashMap; use std::ffi::OsStr; +use std::process::ExitStatus; use std::{ fs, path::{Path, PathBuf}, @@ -12,9 +14,13 @@ use crate::shell::Shell; use indexmap::IndexMap; use rattler_conda_types::Platform; +const ENV_START_SEPERATOR: &str = "<=== RATTLER ENV START ===>"; + /// Type of modification done to the `PATH` variable +#[derive(Default, Clone)] pub enum PathModificationBehaviour { /// Replaces the complete path variable with specified paths. + #[default] Replace, /// Appends the new path variables to the path. E.g. '$PATH:/new/path' Append, @@ -24,6 +30,7 @@ pub enum PathModificationBehaviour { /// A struct that contains the values of the environment variables that are relevant for the activation process. /// The values are stored as strings. Currently, only the `PATH` and `CONDA_PREFIX` environment variables are used. +#[derive(Default, Clone)] pub struct ActivationVariables { /// The value of the `CONDA_PREFIX` environment variable that contains the activated conda prefix path pub conda_prefix: Option, @@ -48,6 +55,7 @@ impl ActivationVariables { /// A struct that holds values for the activation and deactivation /// process of an environment, e.g. activation scripts to execute or environment variables to set. +#[derive(Debug)] pub struct Activator { /// The path to the root of the conda environment pub target_prefix: PathBuf, @@ -138,7 +146,23 @@ pub enum ActivationError { /// An error that occurs when writing the activation script to a file fails #[error("Failed to write activation script to file {0}")] - FailedToWriteActivationScript(#[source] std::fmt::Error), + FailedToWriteActivationScript(#[from] std::fmt::Error), + + /// Failed to run the activation script + #[error("Failed to run activation script (status: {status})")] + FailedToRunActivationScript { + /// The contents of the activation script that was run + script: String, + + /// The stdout output of executing the script + stdout: String, + + /// The stderr output of executing the script + stderr: String, + + /// The error code of running the script + status: ExitStatus, + }, } /// Collect all environment variables that are set in a conda environment. @@ -383,11 +407,74 @@ impl Activator { Ok(ActivationResult { script, path }) } + + /// Runs the activation script and returns the environment variables changed in the environment + /// after running the script. + /// TODO: This only handles UTF-8 formatted strings.. + pub fn run_activation( + &self, + variables: ActivationVariables, + ) -> Result, ActivationError> { + let activation_script = self.activation(variables)?.script; + + // Create a script that starts by emitting all environment variables, then runs the + // activation script followed by again emitting all environment variables. Any changes + // should then become visible. + let mut activation_detection_script = String::new(); + self.shell_type.env(&mut activation_detection_script)?; + self.shell_type + .echo(&mut activation_detection_script, ENV_START_SEPERATOR)?; + activation_detection_script = + format!("{}{}", &activation_detection_script, &activation_script); + self.shell_type + .echo(&mut activation_detection_script, ENV_START_SEPERATOR)?; + self.shell_type.env(&mut activation_detection_script)?; + + // Create a temporary file that we can execute with our shell. + let activation_script_dir = tempfile::TempDir::new()?; + let activation_script_path = activation_script_dir + .path() + .join(format!("activation.{}", self.shell_type.extension())); + fs::write(&activation_script_path, &activation_detection_script)?; + + // Get only the path to the temporary file + let activation_result = self + .shell_type + .create_run_script_command(&activation_script_path) + .output()?; + + if !activation_result.status.success() { + return Err(ActivationError::FailedToRunActivationScript { + script: activation_detection_script, + stdout: String::from_utf8_lossy(&activation_result.stdout).into_owned(), + stderr: String::from_utf8_lossy(&activation_result.stderr).into_owned(), + status: activation_result.status, + }); + } + + let stdout = String::from_utf8_lossy(&activation_result.stdout); + let (before_env, rest) = stdout + .split_once(ENV_START_SEPERATOR) + .unwrap_or(("", stdout.as_ref())); + let (_, after_env) = rest.rsplit_once(ENV_START_SEPERATOR).unwrap_or(("", "")); + + // Parse both environments and find the difference + let before_env = self.shell_type.parse_env(before_env); + let after_env = self.shell_type.parse_env(after_env); + + // Find and return the differences + Ok(after_env + .into_iter() + .filter(|(key, value)| before_env.get(key) != Some(value)) + .map(|(key, value)| (key.to_owned(), value.to_owned())) + .collect()) + } } #[cfg(test)] mod tests { use crate::shell; + use std::collections::BTreeMap; use std::str::FromStr; use super::*; @@ -395,6 +482,7 @@ mod tests { #[cfg(unix)] use crate::activation::PathModificationBehaviour; + use crate::shell::ShellEnum; #[test] fn test_collect_scripts() { @@ -598,4 +686,102 @@ mod tests { let script = get_script(shell::Xonsh, PathModificationBehaviour::Append); insta::assert_snapshot!(script); } + + fn test_run_activation(shell: ShellEnum) { + let environment_dir = tempfile::TempDir::new().unwrap(); + + // Write some environment variables to the `conda-meta/state` folder. + let state_path = environment_dir.path().join("conda-meta/state"); + fs::create_dir_all(state_path.parent().unwrap()).unwrap(); + let quotes = r#"{"env_vars": {"STATE": "Hello, world!"}}"#; + fs::write(&state_path, quotes).unwrap(); + + // Write package specific environment variables + let content_pkg_1 = r#"{"PKG1": "Hello, world!"}"#; + let content_pkg_2 = r#"{"PKG2": "Hello, world!"}"#; + + let env_var_d = environment_dir.path().join("etc/conda/env_vars.d"); + fs::create_dir_all(&env_var_d).expect("Could not create env vars directory"); + + let pkg1 = env_var_d.join("pkg1.json"); + let pkg2 = env_var_d.join("pkg2.json"); + + fs::write(&pkg1, content_pkg_1).expect("could not write file"); + fs::write(&pkg2, content_pkg_2).expect("could not write file"); + + // Write a script that emits a random environment variable via a shell + let mut activation_script = String::new(); + shell + .set_env_var(&mut activation_script, "SCRIPT_ENV", "Hello, world!") + .unwrap(); + + let activation_script_dir = environment_dir.path().join("etc/conda/activate.d"); + fs::create_dir_all(&activation_script_dir).unwrap(); + + fs::write( + activation_script_dir.join(format!("pkg1.{}", shell.extension())), + activation_script, + ) + .unwrap(); + + // Create an activator for the environment + let activator = + Activator::from_path(environment_dir.path(), shell.clone(), Platform::current()) + .unwrap(); + let activation_env = activator + .run_activation(ActivationVariables::default()) + .unwrap(); + + // Diff with the current environment + let current_env = std::env::vars().collect::>(); + let mut env_diff = activation_env + .into_iter() + .filter(|(key, value)| current_env.get(key) != Some(value)) + .collect::>(); + + // Remove system specific environment variables. + env_diff.remove("CONDA_PREFIX"); + env_diff.remove("Path"); + env_diff.remove("PATH"); + + insta::assert_yaml_snapshot!(shell.executable(), env_diff); + } + + #[test] + #[cfg(windows)] + fn test_run_activation_powershell() { + test_run_activation(crate::shell::PowerShell::default().into()) + } + + #[test] + #[cfg(windows)] + fn test_run_activation_cmd() { + test_run_activation(crate::shell::CmdExe::default().into()) + } + + #[test] + #[cfg(unix)] + fn test_run_activation_bash() { + test_run_activation(crate::shell::Bash::default().into()) + } + + #[test] + #[cfg(target_os = "macos")] + fn test_run_activation_zsh() { + test_run_activation(crate::shell::Zsh::default().into()) + } + + #[test] + #[cfg(unix)] + #[ignore] + fn test_run_activation_fish() { + test_run_activation(crate::shell::Fish::default().into()) + } + + #[test] + #[cfg(unix)] + #[ignore] + fn test_run_activation_xonsh() { + test_run_activation(crate::shell::Xonsh::default().into()) + } } diff --git a/crates/rattler_shell/src/shell/mod.rs b/crates/rattler_shell/src/shell/mod.rs index 64da330e9..af989f317 100644 --- a/crates/rattler_shell/src/shell/mod.rs +++ b/crates/rattler_shell/src/shell/mod.rs @@ -4,6 +4,7 @@ use crate::activation::PathModificationBehaviour; use enum_dispatch::enum_dispatch; use itertools::Itertools; use rattler_conda_types::Platform; +use std::collections::HashMap; use std::process::Command; use std::{ fmt::Write, @@ -93,6 +94,23 @@ pub trait Shell { fn format_env_var(&self, var_name: &str) -> String { format!("${{{var_name}}}") } + + /// Emits echoing certain text to stdout. + fn echo(&self, f: &mut impl Write, text: &str) -> std::fmt::Result { + writeln!(f, "echo {}", shlex::quote(text)) + } + + /// Emits writing all current environment variables to stdout. + fn env(&self, f: &mut impl Write) -> std::fmt::Result { + writeln!(f, "/usr/bin/env") + } + + /// Parses environment variables emitted by the `Shell::env` command. + fn parse_env<'i>(&self, env: &'i str) -> HashMap<&'i str, &'i str> { + env.lines() + .filter_map(|line| line.split_once('=')) + .collect() + } } /// Convert a native PATH on Windows to a Unix style path usign cygpath. @@ -123,7 +141,7 @@ fn native_path_to_unix(path: &str) -> Result { } /// A [`Shell`] implementation for the Bash shell. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Default)] pub struct Bash; impl Shell for Bash { @@ -194,7 +212,7 @@ impl Shell for Bash { } /// A [`Shell`] implementation for the Zsh shell. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Default)] pub struct Zsh; impl Shell for Zsh { @@ -226,7 +244,7 @@ impl Shell for Zsh { } /// A [`Shell`] implementation for the Xonsh shell. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Default)] pub struct Xonsh; impl Shell for Xonsh { @@ -258,7 +276,7 @@ impl Shell for Xonsh { } /// A [`Shell`] implementation for the cmd.exe shell. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Default)] pub struct CmdExe; impl Shell for CmdExe { @@ -299,6 +317,15 @@ impl Shell for CmdExe { fn format_env_var(&self, var_name: &str) -> String { format!("%{var_name}%") } + + fn echo(&self, f: &mut impl Write, text: &str) -> std::fmt::Result { + writeln!(f, "@ECHO {}", shlex::quote(text)) + } + + /// Emits writing all current environment variables to stdout. + fn env(&self, f: &mut impl Write) -> std::fmt::Result { + writeln!(f, "@SET") + } } /// A [`Shell`] implementation for PowerShell. @@ -337,10 +364,15 @@ impl Shell for PowerShell { fn format_env_var(&self, var_name: &str) -> String { format!("$Env:{var_name}") } + + /// Emits writing all current environment variables to stdout. + fn env(&self, f: &mut impl Write) -> std::fmt::Result { + writeln!(f, r##"dir env: | %{{"{{0}}={{1}}" -f $_.Name,$_.Value}}"##) + } } /// A [`Shell`] implementation for the Fish shell. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Default)] pub struct Fish; impl Shell for Fish { diff --git a/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__bash.snap b/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__bash.snap new file mode 100644 index 000000000..533fb9420 --- /dev/null +++ b/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__bash.snap @@ -0,0 +1,9 @@ +--- +source: crates/rattler_shell/src/activation.rs +expression: env_diff +--- +PKG1: "Hello, world!" +PKG2: "Hello, world!" +SCRIPT_ENV: "Hello, world!" +STATE: "Hello, world!" + diff --git a/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__cmd.exe.snap b/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__cmd.exe.snap new file mode 100644 index 000000000..533fb9420 --- /dev/null +++ b/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__cmd.exe.snap @@ -0,0 +1,9 @@ +--- +source: crates/rattler_shell/src/activation.rs +expression: env_diff +--- +PKG1: "Hello, world!" +PKG2: "Hello, world!" +SCRIPT_ENV: "Hello, world!" +STATE: "Hello, world!" + diff --git a/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__pwsh.snap b/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__pwsh.snap new file mode 100644 index 000000000..533fb9420 --- /dev/null +++ b/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__pwsh.snap @@ -0,0 +1,9 @@ +--- +source: crates/rattler_shell/src/activation.rs +expression: env_diff +--- +PKG1: "Hello, world!" +PKG2: "Hello, world!" +SCRIPT_ENV: "Hello, world!" +STATE: "Hello, world!" + diff --git a/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__zsh.snap b/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__zsh.snap new file mode 100644 index 000000000..533fb9420 --- /dev/null +++ b/crates/rattler_shell/src/snapshots/rattler_shell__activation__tests__zsh.snap @@ -0,0 +1,9 @@ +--- +source: crates/rattler_shell/src/activation.rs +expression: env_diff +--- +PKG1: "Hello, world!" +PKG2: "Hello, world!" +SCRIPT_ENV: "Hello, world!" +STATE: "Hello, world!" +