diff --git a/crates/uv-cli/src/lib.rs b/crates/uv-cli/src/lib.rs index ce763e0e5..5a18d1c9a 100644 --- a/crates/uv-cli/src/lib.rs +++ b/crates/uv-cli/src/lib.rs @@ -5194,7 +5194,7 @@ pub enum ToolCommand { /// The tool executable directory is determined according to the XDG standard and can be /// retrieved with `uv tool dir --bin`. #[command(alias = "ensurepath")] - UpdateShell, + UpdateShell(PythonUpdateShellArgs), /// Show the path to the uv tools directory. /// /// The tools directory is used to store environments and metadata for installed tools. @@ -6010,7 +6010,16 @@ pub enum PythonCommand { /// The Python executable directory is determined according to the XDG standard and can be /// retrieved with `uv python dir --bin`. #[command(alias = "ensurepath")] - UpdateShell, + UpdateShell(PythonUpdateShellArgs), +} + +#[derive(Args)] +pub struct PythonUpdateShellArgs { + /// Force prepending the executable directory to PATH even if it's already present. + /// + /// This is useful when the directory is in PATH but binaries are shadowed by earlier entries. + #[arg(long)] + pub force: bool, } #[derive(Args)] diff --git a/crates/uv-shell/src/windows.rs b/crates/uv-shell/src/windows.rs index 818d5b999..df0d8a83c 100644 --- a/crates/uv-shell/src/windows.rs +++ b/crates/uv-shell/src/windows.rs @@ -37,7 +37,7 @@ pub fn prepend_path(path: &Path) -> anyhow::Result { } /// Set the windows `PATH` variable in the registry. -fn apply_windows_path_var(path: &HSTRING) -> anyhow::Result<()> { +pub fn apply_windows_path_var(path: &HSTRING) -> anyhow::Result<()> { let environment = CURRENT_USER.create("Environment")?; if path.is_empty() { @@ -52,7 +52,7 @@ fn apply_windows_path_var(path: &HSTRING) -> anyhow::Result<()> { /// Retrieve the windows `PATH` variable from the registry. /// /// Returns `Ok(None)` if the `PATH` variable is not a string. -fn get_windows_path_var() -> anyhow::Result> { +pub fn get_windows_path_var() -> anyhow::Result> { let environment = CURRENT_USER .create("Environment") .context("Failed to open `Environment` key")?; @@ -72,7 +72,7 @@ fn get_windows_path_var() -> anyhow::Result> { /// Prepend a path to the `PATH` variable in the Windows registry. /// /// Returns `Ok(None)` if the given path is already in `PATH`. -fn prepend_to_path(existing_path: &HSTRING, path: HSTRING) -> Option { +pub fn prepend_to_path(existing_path: &HSTRING, path: HSTRING) -> Option { if existing_path.is_empty() { Some(path) } else if existing_path.windows(path.len()).any(|p| *p == *path) { @@ -84,3 +84,60 @@ fn prepend_to_path(existing_path: &HSTRING, path: HSTRING) -> Option { Some(HSTRING::from(new_path)) } } + +/// Force prepend a path to the `PATH` variable, removing any existing entry first. +/// +/// Used with `--force` to ensure the path is at the beginning of PATH. +pub fn force_prepend_path(path: &Path) -> anyhow::Result<()> { + let windows_path = get_windows_path_var()?.unwrap_or_default(); + + // Remove the path if it exists + let new_path = remove_from_path(&windows_path, path); + + // Prepend the path + let final_path = + prepend_to_path(&new_path, HSTRING::from(path)).unwrap_or_else(|| HSTRING::from(path)); + + // Set the `PATH` variable in the registry. + apply_windows_path_var(&final_path)?; + + Ok(()) +} + +/// Check if a path is already in the Windows PATH variable (read-only). +pub fn contains_path(path: &Path) -> anyhow::Result { + let windows_path = get_windows_path_var()?.unwrap_or_default(); + if windows_path.is_empty() { + return Ok(false); + } + + let path_str = path.to_string_lossy(); + let path_string = windows_path.to_string_lossy(); + let is_in_path = path_string + .split(';') + .any(|p| p.trim() == path_str.as_ref()); + + Ok(is_in_path) +} + +/// Remove a path from the `PATH` variable string. +/// +/// Used with `--force` to remove an existing entry before prepending it again. +pub fn remove_from_path(existing_path: &HSTRING, path_to_remove: &Path) -> HSTRING { + if existing_path.is_empty() { + return existing_path.clone(); + } + + let path_str = path_to_remove.to_string_lossy(); + let path_string = existing_path.to_string_lossy(); + let paths: Vec<&str> = path_string + .split(';') + .filter(|p| p.trim() != path_str.as_ref()) + .collect(); + + if paths.is_empty() { + HSTRING::new() + } else { + HSTRING::from(paths.join(";")) + } +} diff --git a/crates/uv/src/commands/python/update_shell.rs b/crates/uv/src/commands/python/update_shell.rs index 18757ff9e..46fb1fd2b 100644 --- a/crates/uv/src/commands/python/update_shell.rs +++ b/crates/uv/src/commands/python/update_shell.rs @@ -7,6 +7,7 @@ use owo_colors::OwoColorize; use tokio::io::AsyncWriteExt; use tracing::debug; +use uv_cli::PythonUpdateShellArgs; use uv_fs::Simplified; use uv_python::managed::python_executable_dir; use uv_shell::Shell; @@ -15,7 +16,10 @@ use crate::commands::ExitStatus; use crate::printer::Printer; /// Ensure that the executable directory is in PATH. -pub(crate) async fn update_shell(printer: Printer) -> Result { +pub(crate) async fn update_shell( + args: PythonUpdateShellArgs, + printer: Printer, +) -> Result { let executable_directory = python_executable_dir()?; debug!( "Ensuring that the executable directory is in PATH: {}", @@ -24,25 +28,44 @@ pub(crate) async fn update_shell(printer: Printer) -> Result { #[cfg(windows)] { - if uv_shell::windows::prepend_path(&executable_directory)? { + // Check if path is in PATH (read-only, doesn't mutate) + let is_in_path = uv_shell::windows::contains_path(&executable_directory)?; + + if is_in_path && !args.force { + // Already in PATH and not forcing - exit early + writeln!( + printer.stderr(), + "Executable directory {} is already in PATH", + executable_directory.simplified_display().cyan() + )?; + return Ok(ExitStatus::Success); + } + + if is_in_path && args.force { + // Already in PATH but forcing - remove it first, then prepend + uv_shell::windows::force_prepend_path(&executable_directory)?; + + writeln!( + printer.stderr(), + "Force updated PATH to prioritize executable directory {}", + executable_directory.simplified_display().cyan() + )?; + writeln!(printer.stderr(), "Restart your shell to apply changes")?; + } else if !is_in_path { + // Not in PATH - normal prepend + uv_shell::windows::prepend_path(&executable_directory)?; writeln!( printer.stderr(), "Updated PATH to include executable directory {}", executable_directory.simplified_display().cyan() )?; writeln!(printer.stderr(), "Restart your shell to apply changes")?; - } else { - writeln!( - printer.stderr(), - "Executable directory {} is already in PATH", - executable_directory.simplified_display().cyan() - )?; } return Ok(ExitStatus::Success); } - if Shell::contains_path(&executable_directory) { + if Shell::contains_path(&executable_directory) && !args.force { writeln!( printer.stderr(), "Executable directory {} is already in PATH", @@ -82,35 +105,70 @@ pub(crate) async fn update_shell(printer: Printer) -> Result { // Search for the command in the file, to avoid redundant updates. match fs_err::tokio::read_to_string(&file).await { Ok(contents) => { - if contents + // Check if command already exists + let command_exists = contents .lines() .map(str::trim) .filter(|line| !line.starts_with('#')) - .any(|line| line.contains(&command)) - { - debug!( - "Skipping already-updated configuration file: {}", - file.simplified_display() - ); - continue; + .any(|line| line.contains(&command)); + + if command_exists { + if args.force { + // With --force: Remove old entry and add it again at the end (ensures first priority) + let new_contents: String = contents + .lines() + .filter(|line| { + let trimmed = line.trim(); + // Remove old # uv comments and lines containing the command + !(trimmed == "# uv" + || (!trimmed.starts_with('#') && trimmed.contains(&command))) + }) + .collect::>() + .join("\n"); + + // Add the command at the end to ensure it's executed last (first in PATH) + let final_contents = format!("{new_contents}\n# uv\n{command}\n"); + + fs_err::tokio::OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .open(&file) + .await? + .write_all(final_contents.as_bytes()) + .await?; + + writeln!( + printer.stderr(), + "Force updated configuration file: {}", + file.simplified_display().cyan() + )?; + updated = true; + } else { + // Without --force: Skip if already exists + debug!( + "Skipping already-updated configuration file: {}", + file.simplified_display() + ); + } + } else { + // Command doesn't exist: Add it normally + fs_err::tokio::OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .open(&file) + .await? + .write_all(format!("{contents}\n# uv\n{command}\n").as_bytes()) + .await?; + + writeln!( + printer.stderr(), + "Updated configuration file: {}", + file.simplified_display().cyan() + )?; + updated = true; } - - // Append the command to the file. - fs_err::tokio::OpenOptions::new() - .create(true) - .truncate(true) - .write(true) - .open(&file) - .await? - .write_all(format!("{contents}\n# uv\n{command}\n").as_bytes()) - .await?; - - writeln!( - printer.stderr(), - "Updated configuration file: {}", - file.simplified_display().cyan() - )?; - updated = true; } Err(err) if err.kind() == std::io::ErrorKind::NotFound => { // Ensure that the directory containing the file exists. diff --git a/crates/uv/src/commands/tool/update_shell.rs b/crates/uv/src/commands/tool/update_shell.rs index ab3b28a93..91fe063b8 100644 --- a/crates/uv/src/commands/tool/update_shell.rs +++ b/crates/uv/src/commands/tool/update_shell.rs @@ -7,6 +7,7 @@ use owo_colors::OwoColorize; use tokio::io::AsyncWriteExt; use tracing::debug; +use uv_cli::PythonUpdateShellArgs; use uv_fs::Simplified; use uv_shell::Shell; use uv_tool::tool_executable_dir; @@ -15,7 +16,10 @@ use crate::commands::ExitStatus; use crate::printer::Printer; /// Ensure that the executable directory is in PATH. -pub(crate) async fn update_shell(printer: Printer) -> Result { +pub(crate) async fn update_shell( + args: PythonUpdateShellArgs, + printer: Printer, +) -> Result { let executable_directory = tool_executable_dir()?; debug!( "Ensuring that the executable directory is in PATH: {}", @@ -24,25 +28,44 @@ pub(crate) async fn update_shell(printer: Printer) -> Result { #[cfg(windows)] { - if uv_shell::windows::prepend_path(&executable_directory)? { + // Check if path is in PATH (read-only, doesn't mutate) + let is_in_path = uv_shell::windows::contains_path(&executable_directory)?; + + if is_in_path && !args.force { + // Already in PATH and not forcing - exit early + writeln!( + printer.stderr(), + "Executable directory {} is already in PATH", + executable_directory.simplified_display().cyan() + )?; + return Ok(ExitStatus::Success); + } + + if is_in_path && args.force { + // Already in PATH but forcing - remove it first, then prepend + uv_shell::windows::force_prepend_path(&executable_directory)?; + + writeln!( + printer.stderr(), + "Force updated PATH to prioritize executable directory {}", + executable_directory.simplified_display().cyan() + )?; + writeln!(printer.stderr(), "Restart your shell to apply changes")?; + } else if !is_in_path { + // Not in PATH - normal prepend + uv_shell::windows::prepend_path(&executable_directory)?; writeln!( printer.stderr(), "Updated PATH to include executable directory {}", executable_directory.simplified_display().cyan() )?; writeln!(printer.stderr(), "Restart your shell to apply changes")?; - } else { - writeln!( - printer.stderr(), - "Executable directory {} is already in PATH", - executable_directory.simplified_display().cyan() - )?; } return Ok(ExitStatus::Success); } - if Shell::contains_path(&executable_directory) { + if Shell::contains_path(&executable_directory) && !args.force { writeln!( printer.stderr(), "Executable directory {} is already in PATH", @@ -82,35 +105,69 @@ pub(crate) async fn update_shell(printer: Printer) -> Result { // Search for the command in the file, to avoid redundant updates. match fs_err::tokio::read_to_string(&file).await { Ok(contents) => { - if contents + // Check if command already exists + let command_exists = contents .lines() .map(str::trim) .filter(|line| !line.starts_with('#')) - .any(|line| line.contains(&command)) - { - debug!( - "Skipping already-updated configuration file: {}", - file.simplified_display() - ); - continue; + .any(|line| line.contains(&command)); + + if command_exists { + if args.force { + // With --force: Remove old entry and add it again at the end + let new_contents: String = contents + .lines() + .filter(|line| { + let trimmed = line.trim(); + // Remove old # uv comments and lines containing the command + !(trimmed == "# uv" + || (!trimmed.starts_with('#') && trimmed.contains(&command))) + }) + .collect::>() + .join("\n"); + + let final_contents = format!("{new_contents}\n# uv\n{command}\n"); + + fs_err::tokio::OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .open(&file) + .await? + .write_all(final_contents.as_bytes()) + .await?; + + writeln!( + printer.stderr(), + "Force updated configuration file: {}", + file.simplified_display().cyan() + )?; + updated = true; + } else { + // Without --force: Skip if already exists + debug!( + "Skipping already-updated configuration file: {}", + file.simplified_display() + ); + } + } else { + // Command doesn't exist: Add it normally + fs_err::tokio::OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .open(&file) + .await? + .write_all(format!("{contents}\n# uv\n{command}\n").as_bytes()) + .await?; + + writeln!( + printer.stderr(), + "Updated configuration file: {}", + file.simplified_display().cyan() + )?; + updated = true; } - - // Append the command to the file. - fs_err::tokio::OpenOptions::new() - .create(true) - .truncate(true) - .write(true) - .open(&file) - .await? - .write_all(format!("{contents}\n# uv\n{command}\n").as_bytes()) - .await?; - - writeln!( - printer.stderr(), - "Updated configuration file: {}", - file.simplified_display().cyan() - )?; - updated = true; } Err(err) if err.kind() == std::io::ErrorKind::NotFound => { // Ensure that the directory containing the file exists. diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs index d8d4dbb91..9e1b85bb0 100644 --- a/crates/uv/src/lib.rs +++ b/crates/uv/src/lib.rs @@ -1551,9 +1551,9 @@ async fn run(mut cli: Cli) -> Result { commands::tool_uninstall(args.name, printer).await } Commands::Tool(ToolNamespace { - command: ToolCommand::UpdateShell, + command: ToolCommand::UpdateShell(args), }) => { - commands::tool_update_shell(printer).await?; + commands::tool_update_shell(args, printer).await?; Ok(ExitStatus::Success) } Commands::Tool(ToolNamespace { @@ -1745,9 +1745,9 @@ async fn run(mut cli: Cli) -> Result { Ok(ExitStatus::Success) } Commands::Python(PythonNamespace { - command: PythonCommand::UpdateShell, + command: PythonCommand::UpdateShell(args), }) => { - commands::python_update_shell(printer).await?; + commands::python_update_shell(args, printer).await?; Ok(ExitStatus::Success) } Commands::Publish(args) => { diff --git a/crates/uv/tests/it/common/mod.rs b/crates/uv/tests/it/common/mod.rs index 46ea71091..3010c01ad 100644 --- a/crates/uv/tests/it/common/mod.rs +++ b/crates/uv/tests/it/common/mod.rs @@ -1289,6 +1289,14 @@ impl TestContext { command } + /// Create a `uv python update-shell` command with options shared across scenarios. + pub fn python_update_shell(&self) -> Command { + let mut command = Self::new_command(); + command.arg("python").arg("update-shell"); + self.add_shared_options(&mut command, true); + command + } + /// Create a `uv run` command with options shared across scenarios. pub fn run(&self) -> Command { let mut command = Self::new_command(); diff --git a/crates/uv/tests/it/main.rs b/crates/uv/tests/it/main.rs index b81292b49..ade100216 100644 --- a/crates/uv/tests/it/main.rs +++ b/crates/uv/tests/it/main.rs @@ -102,6 +102,9 @@ mod python_pin; #[cfg(feature = "python-managed")] mod python_upgrade; +#[cfg(feature = "python")] +mod python_update_shell; + #[cfg(all(feature = "python", feature = "pypi"))] mod run; diff --git a/crates/uv/tests/it/python_update_shell.rs b/crates/uv/tests/it/python_update_shell.rs new file mode 100644 index 000000000..858af3d7a --- /dev/null +++ b/crates/uv/tests/it/python_update_shell.rs @@ -0,0 +1,176 @@ +use assert_cmd::assert::OutputAssertExt; +use assert_fs::fixture::PathChild; + +use uv_static::EnvVars; + +use crate::common::{TestContext, uv_snapshot}; + +#[test] +fn python_update_shell_not_in_path() { + let context = TestContext::new("3.12"); + + #[cfg(not(windows))] + { + // Zsh uses .zshenv, not .zshrc + let shell_config = context.home_dir.child(".zshenv"); + + uv_snapshot!(context.filters(), context + .python_update_shell() + .env(EnvVars::HOME, context.home_dir.as_os_str()) + .env("SHELL", "/bin/zsh"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Created configuration file: [HOME]/.zshenv + Restart your shell to apply changes + "###); + + // Verify the file was created with the correct content + let contents = fs_err::read_to_string(shell_config.path()).unwrap(); + assert!(contents.contains("export PATH=")); + assert!(contents.contains("# uv")); + } + + #[cfg(windows)] + { + // On Windows, PATH is updated in the registry, not a config file + uv_snapshot!(context.filters(), context + .python_update_shell() + .env(EnvVars::HOME, context.home_dir.as_os_str()), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Updated PATH to include executable directory [USER_CONFIG_DIR]/data/../bin + Restart your shell to apply changes + "###); + } +} + +#[test] +fn python_update_shell_already_in_path() { + let context = TestContext::new("3.12"); + + // Set a specific bin directory using UV_PYTHON_BIN_DIR + let bin_dir = context.home_dir.child("bin"); + fs_err::create_dir_all(bin_dir.path()).unwrap(); + + #[cfg(not(windows))] + { + // Set PATH to include the bin directory so it's "already in PATH" + let path_with_bin = + std::env::join_paths(std::iter::once(bin_dir.path().to_path_buf()).chain( + std::env::split_paths(&std::env::var(EnvVars::PATH).unwrap_or_default()), + )) + .unwrap(); + + // Run without --force - should skip because it's already in PATH + uv_snapshot!(context.filters(), context + .python_update_shell() + .env(EnvVars::HOME, context.home_dir.as_os_str()) + .env(EnvVars::UV_PYTHON_BIN_DIR, bin_dir.path()) + .env(EnvVars::PATH, path_with_bin) + .env("SHELL", "/bin/zsh"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Executable directory [HOME]/bin is already in PATH + "###); + } + + #[cfg(windows)] + { + // On Windows, contains_path checks the registry, not env vars + // Since we can't easily set up the registry in tests, we'll test + // that the command succeeds. The "already in PATH" check will + // depend on the actual registry state, which we can't control. + // This test verifies the command works even if the path is already there. + context + .python_update_shell() + .env(EnvVars::HOME, context.home_dir.as_os_str()) + .env(EnvVars::UV_PYTHON_BIN_DIR, bin_dir.path()) + .assert() + .success(); + } +} + +#[test] +fn python_update_shell_force() { + let context = TestContext::new("3.12"); + + #[cfg(not(windows))] + { + // Zsh uses .zshenv, not .zshrc + let shell_config = context.home_dir.child(".zshenv"); + + // First run - add to PATH + context + .python_update_shell() + .env(EnvVars::HOME, context.home_dir.as_os_str()) + .env("SHELL", "/bin/zsh") + .assert() + .success(); + + let first_contents = fs_err::read_to_string(shell_config.path()).unwrap(); + let first_count = first_contents.matches("export PATH=").count(); + + // Second run with --force - should update + uv_snapshot!(context.filters(), context + .python_update_shell() + .arg("--force") + .env(EnvVars::HOME, context.home_dir.as_os_str()) + .env("SHELL", "/bin/zsh"), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Force updated configuration file: [HOME]/.zshenv + Restart your shell to apply changes + "###); + + // Verify only one PATH export exists (old one removed, new one added) + let second_contents = fs_err::read_to_string(shell_config.path()).unwrap(); + let second_count = second_contents.matches("export PATH=").count(); + assert_eq!( + first_count, second_count, + "Should have same number of PATH exports" + ); + + // Verify the command is at the end + assert!( + second_contents.trim_end().ends_with("export PATH=") + || second_contents.trim_end().ends_with('"') + ); + } + + #[cfg(windows)] + { + // On Windows, --force updates the registry + // First run - add to PATH + context + .python_update_shell() + .env(EnvVars::HOME, context.home_dir.as_os_str()) + .assert() + .success(); + + // Second run with --force - should update registry + uv_snapshot!(context.filters(), context + .python_update_shell() + .arg("--force") + .env(EnvVars::HOME, context.home_dir.as_os_str()), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Force updated PATH to prioritize executable directory [USER_CONFIG_DIR]/data/../bin + Restart your shell to apply changes + "###); + } +}