Skip to content

Commit

Permalink
Treat the base Conda environment as a system environment (#7691)
Browse files Browse the repository at this point in the history
Closes #7124
Closes #7137
  • Loading branch information
zanieb committed Oct 24, 2024
1 parent d9d55b7 commit 248e9f7
Show file tree
Hide file tree
Showing 4 changed files with 152 additions and 23 deletions.
51 changes: 34 additions & 17 deletions crates/uv-python/src/discovery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};

Expand Down Expand Up @@ -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`
Expand Down Expand Up @@ -227,27 +229,27 @@ 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<Item = Result<(PythonSource, PathBuf), Error>> + '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)
.map(|path| Ok((PythonSource::ActiveEnvironment, path)))
})
.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)))
Expand All @@ -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)
}
Expand Down Expand Up @@ -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))
}
}
}

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -1826,6 +1842,7 @@ impl VersionRequest {
Self::Default => match source {
PythonSource::ParentInterpreter
| PythonSource::CondaPrefix
| PythonSource::BaseCondaPrefix
| PythonSource::ProvidedPath
| PythonSource::DiscoveredEnvironment
| PythonSource::ActiveEnvironment => Self::Any,
Expand Down Expand Up @@ -2217,7 +2234,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"),
Expand Down
68 changes: 68 additions & 0 deletions crates/uv-python/src/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(())
}

Expand Down
53 changes: 47 additions & 6 deletions crates/uv-python/src/virtualenv.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,15 +57,56 @@ pub(crate) fn virtualenv_from_env() -> Option<PathBuf> {
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/<name>`.
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<PathBuf> {
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<PathBuf> {
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.
Expand Down
3 changes: 3 additions & 0 deletions crates/uv-static/src/env_vars.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand Down

0 comments on commit 248e9f7

Please sign in to comment.