feat: add support for PIXI_ENVIRONMENT_NAME and PS1 prompt modification

docs/reference/cli/pixi/exec.md

@@ -33,6 +33,9 @@ pixi exec [OPTIONS] [COMMAND]...
 :  If specified a new environment is always created even if one already exists
 - <a id="arg---list" href="#arg---list">`--list <LIST>`</a>
 :  Before executing the command, list packages in the environment Specify `--list=some_regex` to filter the shown packages
+- <a id="arg---no-modify-ps1" href="#arg---no-modify-ps1">`--no-modify-ps1`</a>
+:  Disable modification of the PS1 prompt to indicate the temporary environment
+<br>**default**: `false`
 
 ## Config Options
 - <a id="arg---auth-file" href="#arg---auth-file">`--auth-file <AUTH_FILE>`</a>

src/cli/exec.rs

@@ -1,4 +1,4 @@
-use std::{path::Path, str::FromStr, sync::LazyLock};
+use std::{collections::BTreeSet, collections::HashMap, path::Path, str::FromStr, sync::LazyLock};
 
 use clap::{Parser, ValueHint};
 use itertools::Itertools;
@@ -12,7 +12,7 @@ use rattler::{
 };
 use rattler_conda_types::{GenericVirtualPackage, MatchSpec, PackageName, Platform};
 use rattler_solve::{SolverImpl, SolverTask, resolvo::Solver};
-use rattler_virtual_packages::{VirtualPackage, VirtualPackageOverrides};
+use rattler_virtual_packages::{VirtualPackageOverrides, VirtualPackages};
 use reqwest_middleware::ClientWithMiddleware;
 use uv_configuration::RAYON_INITIALIZE;
 
@@ -59,6 +59,10 @@ pub struct Args {
     #[clap(long = "list", num_args = 0..=1, default_missing_value = "", require_equals = true)]
     pub list: Option<String>,
 
+    /// Disable modification of the PS1 prompt to indicate the temporary environment
+    #[clap(long, action = clap::ArgAction::SetTrue, default_value = "false")]
+    pub no_modify_ps1: bool,
+
     #[clap(flatten)]
     pub config: ConfigCli,
 }
@@ -68,23 +72,84 @@ pub async fn execute(args: Args) -> miette::Result<()> {
     let config = Config::with_cli_config(&args.config);
     let cache_dir = pixi_config::get_cache_dir().context("failed to determine cache directory")?;
 
-    let mut command_args = args.command.iter();
-    let command = command_args.next().ok_or_else(|| miette::miette!(help ="i.e when specifying specs explicitly use a command at the end: `pixi exec -s python==3.12 python`", "missing required command to execute",))?;
+    let mut command_iter = args.command.iter();
+    let command = command_iter.next().ok_or_else(|| miette::miette!(help ="i.e when specifying specs explicitly use a command at the end: `pixi exec -s python==3.12 python`", "missing required command to execute",))?;
     let (_, client) = build_reqwest_clients(Some(&config), None)?;
 
+    // Determine the specs for installation and for the environment name.
+    let mut name_specs = args.specs.clone();
+    name_specs.extend(args.with.clone());
+
+    let mut install_specs = name_specs.clone();
+
+    // Guess a package from the command if no specs were provided at all OR if --with is used
+    let should_guess_package = name_specs.is_empty() || !args.with.is_empty();
+    if should_guess_package {
+        install_specs.push(guess_package_spec(command));
+    }
+
     // Create the environment to run the command in.
-    let prefix = create_exec_prefix(&args, &cache_dir, &config, &client).await?;
+    let prefix = create_exec_prefix(
+        &args,
+        &install_specs,
+        &cache_dir,
+        &config,
+        &client,
+        should_guess_package,
+    )
+    .await?;
 
     // Get environment variables from the activation
-    let activation_env = run_activation(&prefix).await?;
+    let mut activation_env = run_activation(&prefix).await?;
+
+    // Collect unique package names for environment naming
+    let package_names: BTreeSet<String> = name_specs
+        .iter()
+        .filter_map(|spec| spec.name.as_ref().map(|n| n.as_normalized().to_string()))
+        .collect();
+
+    if !package_names.is_empty() {
+        let env_name = format!(
+            "temp:{}",
+            package_names.into_iter().format(",")
+        );
+
+        activation_env.insert("PIXI_ENVIRONMENT_NAME".into(), env_name.clone());
+
+        if !args.no_modify_ps1 && std::env::current_dir().is_ok() {
+            let (prompt_var, prompt_value) = if cfg!(windows) {
+                ("_PIXI_PROMPT", format!("(pixi:{}) $P$G", env_name))
+            } else {
+                ("PS1", format!(r"(pixi:{}) [\w] \$", env_name))
+            };
+
+            activation_env.insert(prompt_var.into(), prompt_value);
+
+            if cfg!(windows) {
+                activation_env.insert("PROMPT".into(), String::from("$P$G"));
+            }
+        }
+    }
 
     // Ignore CTRL+C so that the child is responsible for its own signal handling.
     let _ctrl_c = tokio::spawn(async { while tokio::signal::ctrl_c().await.is_ok() {} });
 
     // Spawn the command
-    let status = std::process::Command::new(command)
-        .args(command_args)
-        .envs(activation_env.iter().map(|(k, v)| (k.as_str(), v.as_str())))
+    let mut cmd = std::process::Command::new(command);
+    let command_args_vec: Vec<_> = command_iter.collect();
+    cmd.args(&command_args_vec);
+
+    // On Windows, when using cmd.exe or cmd, we need to set environment variables in a special way
+    // because cmd.exe has its own environment variable expansion rules
+    if cfg!(windows) && (command.to_lowercase().ends_with("cmd.exe") || command == "cmd") {
+        let mut env = std::env::vars().collect::<HashMap<String, String>>();
+        env.extend(activation_env);
+        cmd.envs(env);
+    } else {
+        cmd.envs(activation_env.iter().map(|(k, v)| (k.as_str(), v.as_str())));
+    }
+
+    let status = cmd
         .status()
         .into_diagnostic()
         .with_context(|| format!("failed to execute '{}'", &command))?;
@@ -96,20 +161,24 @@ pub async fn execute(args: Args) -> miette::Result<()> {
 /// Creates a prefix for the `pixi exec` command.
 pub async fn create_exec_prefix(
     args: &Args,
+    specs: &[MatchSpec],
     cache_dir: &Path,
     config: &Config,
     client: &ClientWithMiddleware,
+    has_guessed_package: bool,
 ) -> miette::Result<Prefix> {
     let command = args.command.first().expect("missing required command");
-    let specs = args.specs.clone();
+    let specs = specs.to_vec();
+
     let channels = args
         .channels
         .resolve_from_config(config)?
         .iter()
         .map(|c| c.base_url.to_string())
         .collect();
 
-    let environment_hash = EnvironmentHash::new(command.clone(), specs, channels, args.platform);
+    let environment_hash =
+        EnvironmentHash::new(command.clone(), specs.clone(), channels, args.platform);
 
     let prefix = Prefix::new(
         cache_dir
@@ -148,23 +217,6 @@ pub async fn create_exec_prefix(
     // Construct a gateway to get repodata.
     let gateway = config.gateway().with_client(client.clone()).finish();
 
-    // Determine the specs to use for the environment
-    let specs = if args.specs.is_empty() {
-        let command = args.command.first().expect("missing required command");
-        let guessed_spec = guess_package_spec(command);
-
-        tracing::debug!(
-            "no specs provided, guessed {} from command {command}",
-            guessed_spec
-        );
-
-        let mut with_specs = args.with.clone();
-        with_specs.push(guessed_spec);
-        with_specs
-    } else {
-        args.specs.clone()
-    };
-
     let channels = args.channels.resolve_from_config(config)?;
 
     // Get the repodata for the specs
@@ -180,13 +232,12 @@ pub async fn create_exec_prefix(
     .context("failed to get repodata")?;
 
     // Determine virtual packages of the current platform
-    let virtual_packages = VirtualPackage::detect(&VirtualPackageOverrides::from_env())
-        .into_diagnostic()
-        .context("failed to determine virtual packages")?
-        .iter()
-        .cloned()
-        .map(GenericVirtualPackage::from)
-        .collect();
+    let virtual_packages: Vec<GenericVirtualPackage> =
+        VirtualPackages::detect(&VirtualPackageOverrides::from_env())
+            .into_diagnostic()
+            .context("failed to determine virtual packages")?
+            .into_generic_virtual_packages()
+            .collect();
 
     // Solve the environment
     tracing::info!(
@@ -196,16 +247,46 @@ pub async fn create_exec_prefix(
             .unwrap_or(prefix.root())
             .display()
     );
-    let specs_clone = specs.clone();
-    let solved_records = wrap_in_progress("solving environment", move || {
+    let solve_result = wrap_in_progress("solving environment", || {
         Solver.solve(SolverTask {
-            specs: specs_clone,
-            virtual_packages,
-            ..SolverTask::from_iter(&repodata)
+            specs: specs.clone(),
+            virtual_packages: virtual_packages.clone(),
+            ..SolverTask::from_iter(&repodata.clone())
         })
-    })
-    .into_diagnostic()
-    .context("failed to solve environment")?;
+    });
+
+    let (solved_records, final_specs) = match solve_result {
+        Ok(records) => (records, specs.to_vec()),
+        Err(err) if has_guessed_package && !args.with.is_empty() => {
+            // If solving failed and we guessed a package while using --with,
+            // try again without the guessed package (last spec)
+            let guessed_package_name = specs[specs.len() - 1]
+                .name
+                .as_ref()
+                .map(|name| name.as_source())
+                .unwrap_or("<unknown>");
+            tracing::debug!(
+                "Solver failed with guessed package '{}', retrying without it: {}",
+                guessed_package_name,
+                err
+            );
+            let records = wrap_in_progress("retrying solve without guessed package", || {
+                Solver.solve(SolverTask {
+                    specs: specs[..specs.len() - 1].to_vec(),
+                    virtual_packages: virtual_packages.clone(),
+                    ..SolverTask::from_iter(&repodata.clone())
+                })
+            })
+            .into_diagnostic()
+            .context("failed to solve environment even without guessed package")?;
+            (records, specs[..specs.len() - 1].to_vec())
+        }
+        Err(err) => {
+            return Err(err)
+                .into_diagnostic()
+                .context("failed to solve environment");
+        }
+    };
 
     // Force the initialization of the rayon thread pool to avoid implicit creation
     // by the Installer.
@@ -232,7 +313,7 @@ pub async fn create_exec_prefix(
     write_guard.finish().await.into_diagnostic()?;
 
     if let Some(ref regex) = args.list {
-        list_exec_environment(specs, solved_records, regex.clone())?;
+        list_exec_environment(final_specs, solved_records, regex.clone())?;
     }
 
     Ok(prefix)

tests/integration_python/common.py

@@ -4,7 +4,8 @@
 from contextlib import contextmanager
 from enum import IntEnum
 from pathlib import Path
-from typing import Generator
+import sys
+from typing import Generator, Optional, Sequence, Tuple
 
 from rattler import Platform
 
@@ -33,12 +34,12 @@ class ExitCode(IntEnum):
 
 
 class Output:
-    command: list[Path | str]
+    command: Sequence[Path | str]
     stdout: str
     stderr: str
     returncode: int
 
-    def __init__(self, command: list[Path | str], stdout: str, stderr: str, returncode: int):
+    def __init__(self, command: Sequence[Path | str], stdout: str, stderr: str, returncode: int):
         self.command = command
         self.stdout = stdout
         self.stderr = stderr
@@ -49,7 +50,7 @@ def __str__(self) -> str:
 
 
 def verify_cli_command(
-    command: list[Path | str],
+    command: Sequence[Path | str],
     expected_exit_code: ExitCode = ExitCode.SUCCESS,
     stdout_contains: str | list[str] | None = None,
     stdout_excludes: str | list[str] | None = None,
@@ -167,3 +168,34 @@ def cwd(path: str | Path) -> Generator[None, None, None]:
         yield
     finally:
         os.chdir(oldpwd)
+
+
+def run_and_get_env(pixi: Path, *args: str, env_var: str) -> Tuple[Optional[str], Output]:
+    if sys.platform.startswith("win"):
+        cmd = [str(pixi), "exec", *args, "--", "cmd", "/c", f"echo %{env_var}%"]
+    else:
+        cmd = [str(pixi), "exec", *args, "--", "printenv", env_var]
+
+    try:
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+        )
+
+        value = result.stdout.strip()
+
+        output = Output(
+            command=cmd,
+            stdout=value,
+            stderr=result.stderr.strip(),
+            returncode=result.returncode,
+        )
+
+        return (value if value and value != f"%{env_var}%" else None, output)
+    except Exception as e:
+        print(f"Error running command: {e}")
+        print(f"Command: {' '.join(cmd)}")
+        raise