Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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