From 788a034861eb7019dbd7587a75ec7de7e563f7d7 Mon Sep 17 00:00:00 2001 From: Zanie Blue Date: Mon, 21 Oct 2024 17:23:49 -0500 Subject: [PATCH] Treat the base Conda environment as a system environment (#7691) Closes https://github.com/astral-sh/uv/issues/7124 Closes https://github.com/astral-sh/uv/issues/7137 --- crates/uv-python/src/discovery.rs | 51 ++++++++++++++-------- crates/uv-python/src/tests.rs | 68 ++++++++++++++++++++++++++++++ crates/uv-python/src/virtualenv.rs | 53 ++++++++++++++++++++--- crates/uv-static/src/env_vars.rs | 3 ++ 4 files changed, 152 insertions(+), 23 deletions(-) diff --git a/crates/uv-python/src/discovery.rs b/crates/uv-python/src/discovery.rs index 0aa66dc71718..cfee35e8adcf 100644 --- a/crates/uv-python/src/discovery.rs +++ b/crates/uv-python/src/discovery.rs @@ -26,8 +26,8 @@ use crate::microsoft_store::find_microsoft_store_pythons; #[cfg(windows)] use crate::py_launcher::{registry_pythons, WindowsPython}; use crate::virtualenv::{ - conda_prefix_from_env, virtualenv_from_env, virtualenv_from_working_dir, - virtualenv_python_executable, + conda_environment_from_env, virtualenv_from_env, virtualenv_from_working_dir, + virtualenv_python_executable, CondaEnvironmentKind, }; use crate::{Interpreter, PythonVersion}; @@ -179,6 +179,8 @@ pub enum PythonSource { ActiveEnvironment, /// A conda environment was active e.g. via `CONDA_PREFIX` CondaPrefix, + /// A base conda environment was active e.g. via `CONDA_PREFIX` + BaseCondaPrefix, /// An environment was discovered e.g. via `.venv` DiscoveredEnvironment, /// An executable was found in the search path i.e. `PATH` @@ -227,18 +229,17 @@ pub enum Error { SourceNotAllowed(PythonRequest, PythonSource, PythonPreference), } -/// Lazily iterate over Python executables in mutable environments. +/// Lazily iterate over Python executables in mutable virtual environments. /// /// The following sources are supported: /// /// - Active virtual environment (via `VIRTUAL_ENV`) -/// - Active conda environment (via `CONDA_PREFIX`) /// - Discovered virtual environment (e.g. `.venv` in a parent directory) /// /// Notably, "system" environments are excluded. See [`python_executables_from_installed`]. -fn python_executables_from_environments<'a>( +fn python_executables_from_virtual_environments<'a>( ) -> impl Iterator> + 'a { - let from_virtual_environment = std::iter::once_with(|| { + let from_active_environment = std::iter::once_with(|| { virtualenv_from_env() .into_iter() .map(virtualenv_python_executable) @@ -246,8 +247,9 @@ fn python_executables_from_environments<'a>( }) .flatten(); + // N.B. we prefer the conda environment over discovered virtual environments let from_conda_environment = std::iter::once_with(|| { - conda_prefix_from_env() + conda_environment_from_env(CondaEnvironmentKind::Child) .into_iter() .map(virtualenv_python_executable) .map(|path| Ok((PythonSource::CondaPrefix, path))) @@ -265,7 +267,7 @@ fn python_executables_from_environments<'a>( }) .flatten_ok(); - from_virtual_environment + from_active_environment .chain(from_conda_environment) .chain(from_discovered_environment) } @@ -400,23 +402,35 @@ fn python_executables<'a>( }) .flatten(); - let from_environments = python_executables_from_environments(); + // Check if the the base conda environment is active + let from_base_conda_environment = std::iter::once_with(|| { + conda_environment_from_env(CondaEnvironmentKind::Base) + .into_iter() + .map(virtualenv_python_executable) + .map(|path| Ok((PythonSource::BaseCondaPrefix, path))) + }) + .flatten(); + + let from_virtual_environments = python_executables_from_virtual_environments(); let from_installed = python_executables_from_installed(version, implementation, preference); // Limit the search to the relevant environment preference; we later validate that they match // the preference but queries are expensive and we query less interpreters this way. match environments { EnvironmentPreference::OnlyVirtual => { - Box::new(from_parent_interpreter.chain(from_environments)) + Box::new(from_parent_interpreter.chain(from_virtual_environments)) } EnvironmentPreference::ExplicitSystem | EnvironmentPreference::Any => Box::new( from_parent_interpreter - .chain(from_environments) + .chain(from_virtual_environments) + .chain(from_base_conda_environment) + .chain(from_installed), + ), + EnvironmentPreference::OnlySystem => Box::new( + from_parent_interpreter + .chain(from_base_conda_environment) .chain(from_installed), ), - EnvironmentPreference::OnlySystem => { - Box::new(from_parent_interpreter.chain(from_installed)) - } } } @@ -611,8 +625,8 @@ fn satisfies_environment_preference( ) -> bool { match ( preference, - // Conda environments are not conformant virtual environments but we treat them as such - interpreter.is_virtualenv() || matches!(source, PythonSource::CondaPrefix), + // Conda environments are not conformant virtual environments but we treat them as such. + interpreter.is_virtualenv() || (matches!(source, PythonSource::CondaPrefix)), ) { (EnvironmentPreference::Any, _) => true, (EnvironmentPreference::OnlyVirtual, true) => true, @@ -1493,6 +1507,7 @@ impl PythonSource { Self::Managed | Self::Registry | Self::MicrosoftStore => false, Self::SearchPath | Self::CondaPrefix + | Self::BaseCondaPrefix | Self::ProvidedPath | Self::ParentInterpreter | Self::ActiveEnvironment @@ -1505,6 +1520,7 @@ impl PythonSource { match self { Self::Managed | Self::Registry | Self::SearchPath | Self::MicrosoftStore => false, Self::CondaPrefix + | Self::BaseCondaPrefix | Self::ProvidedPath | Self::ParentInterpreter | Self::ActiveEnvironment @@ -1826,6 +1842,7 @@ impl VersionRequest { Self::Default => match source { PythonSource::ParentInterpreter | PythonSource::CondaPrefix + | PythonSource::BaseCondaPrefix | PythonSource::ProvidedPath | PythonSource::DiscoveredEnvironment | PythonSource::ActiveEnvironment => Self::Any, @@ -2236,7 +2253,7 @@ impl fmt::Display for PythonSource { match self { Self::ProvidedPath => f.write_str("provided path"), Self::ActiveEnvironment => f.write_str("active virtual environment"), - Self::CondaPrefix => f.write_str("conda prefix"), + Self::CondaPrefix | Self::BaseCondaPrefix => f.write_str("conda prefix"), Self::DiscoveredEnvironment => f.write_str("virtual environment"), Self::SearchPath => f.write_str("search path"), Self::Registry => f.write_str("registry"), diff --git a/crates/uv-python/src/tests.rs b/crates/uv-python/src/tests.rs index ad5f2d495b1c..2c8a6475f742 100644 --- a/crates/uv-python/src/tests.rs +++ b/crates/uv-python/src/tests.rs @@ -949,6 +949,74 @@ fn find_python_from_conda_prefix() -> Result<()> { "We should allow the active conda python" ); + let baseenv = context.tempdir.child("base"); + TestContext::mock_conda_prefix(&baseenv, "3.12.1")?; + + // But not if it's a base environment + let result = context.run_with_vars( + &[ + ("CONDA_PREFIX", Some(baseenv.as_os_str())), + ("CONDA_DEFAULT_ENV", Some(&OsString::from("base"))), + ], + || { + find_python_installation( + &PythonRequest::Default, + EnvironmentPreference::OnlyVirtual, + PythonPreference::OnlySystem, + &context.cache, + ) + }, + )?; + + assert!( + matches!(result, Err(PythonNotFound { .. })), + "We should not allow the non-virtual environment; got {result:?}" + ); + + // Unless, system interpreters are included... + let python = context.run_with_vars( + &[ + ("CONDA_PREFIX", Some(baseenv.as_os_str())), + ("CONDA_DEFAULT_ENV", Some(&OsString::from("base"))), + ], + || { + find_python_installation( + &PythonRequest::Default, + EnvironmentPreference::OnlySystem, + PythonPreference::OnlySystem, + &context.cache, + ) + }, + )??; + + assert_eq!( + python.interpreter().python_full_version().to_string(), + "3.12.1", + "We should find the base conda environment" + ); + + // If the environment name doesn't match the default, we should not treat it as system + let python = context.run_with_vars( + &[ + ("CONDA_PREFIX", Some(condaenv.as_os_str())), + ("CONDA_DEFAULT_ENV", Some(&OsString::from("base"))), + ], + || { + find_python_installation( + &PythonRequest::Default, + EnvironmentPreference::OnlyVirtual, + PythonPreference::OnlySystem, + &context.cache, + ) + }, + )??; + + assert_eq!( + python.interpreter().python_full_version().to_string(), + "3.12.0", + "We should find the conda environment" + ); + Ok(()) } diff --git a/crates/uv-python/src/virtualenv.rs b/crates/uv-python/src/virtualenv.rs index 865aabf2096a..5fdfc094f029 100644 --- a/crates/uv-python/src/virtualenv.rs +++ b/crates/uv-python/src/virtualenv.rs @@ -57,15 +57,56 @@ pub(crate) fn virtualenv_from_env() -> Option { None } +#[derive(Debug, PartialEq, Eq, Copy, Clone)] +pub(crate) enum CondaEnvironmentKind { + /// The base Conda environment; treated like a system Python environment. + Base, + /// Any other Conda environment; treated like a virtual environment. + Child, +} + +impl CondaEnvironmentKind { + /// Whether the given `CONDA_PREFIX` path is the base Conda environment. + /// + /// When the base environment is used, `CONDA_DEFAULT_ENV` will be set to a name, i.e., `base` or + /// `root` which does not match the prefix, e.g. `/usr/local` instead of + /// `/usr/local/conda/envs/`. + fn from_prefix_path(path: &Path) -> Self { + // If we cannot read `CONDA_DEFAULT_ENV`, there's no way to know if the base environment + let Ok(default_env) = env::var(EnvVars::CONDA_DEFAULT_ENV) else { + return CondaEnvironmentKind::Child; + }; + + // These are the expected names for the base environment + if default_env != "base" && default_env != "root" { + return CondaEnvironmentKind::Child; + } + + let Some(name) = path.file_name() else { + return CondaEnvironmentKind::Child; + }; + + if name.to_str().is_some_and(|name| name == default_env) { + CondaEnvironmentKind::Base + } else { + CondaEnvironmentKind::Child + } + } +} + /// Locate an active conda environment by inspecting environment variables. /// -/// Supports `CONDA_PREFIX`. -pub(crate) fn conda_prefix_from_env() -> Option { - if let Some(dir) = env::var_os(EnvVars::CONDA_PREFIX).filter(|value| !value.is_empty()) { - return Some(PathBuf::from(dir)); - } +/// If `base` is true, the active environment must be the base environment or `None` is returned, +/// and vice-versa. +pub(crate) fn conda_environment_from_env(kind: CondaEnvironmentKind) -> Option { + let dir = env::var_os(EnvVars::CONDA_PREFIX).filter(|value| !value.is_empty())?; + let path = PathBuf::from(dir); - None + if kind != CondaEnvironmentKind::from_prefix_path(&path) { + return None; + }; + + Some(path) } /// Locate a virtual environment by searching the file system. diff --git a/crates/uv-static/src/env_vars.rs b/crates/uv-static/src/env_vars.rs index 8e59f6de1eb7..de642bc6014a 100644 --- a/crates/uv-static/src/env_vars.rs +++ b/crates/uv-static/src/env_vars.rs @@ -247,6 +247,9 @@ impl EnvVars { /// Used to detect an activated Conda environment. pub const CONDA_PREFIX: &'static str = "CONDA_PREFIX"; + /// Used to determine if an active Conda environment is the base environment or not. + pub const CONDA_DEFAULT_ENV: &'static str = "CONDA_DEFAULT_ENV"; + /// Disables prepending virtual environment name to the terminal prompt. pub const VIRTUAL_ENV_DISABLE_PROMPT: &'static str = "VIRTUAL_ENV_DISABLE_PROMPT";