Treat the base Conda environment as a system environment (#7691)

zanieb · zanieb · commit 8034de7295b9 · 2024-11-07T14:29:54.000-06:00
Closes #7124 Closes #7137
diff --git 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};
 
@@ -185,6 +185,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`
@@ -233,27 +235,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)))
@@ -271,7 +273,7 @@ fn python_executables_from_environments<'a>(
     })
     .flatten_ok();
 
-    from_virtual_environment
+    from_active_environment
         .chain(from_conda_environment)
         .chain(from_discovered_environment)
 }
@@ -406,23 +408,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))
-        }
     }
 }
 
@@ -617,8 +631,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,
@@ -1515,6 +1529,7 @@ impl PythonSource {
             Self::Managed | Self::Registry | Self::MicrosoftStore => false,
             Self::SearchPath
             | Self::CondaPrefix
+            | Self::BaseCondaPrefix
             | Self::ProvidedPath
             | Self::ParentInterpreter
             | Self::ActiveEnvironment
@@ -1527,6 +1542,7 @@ impl PythonSource {
         match self {
             Self::Managed | Self::Registry | Self::SearchPath | Self::MicrosoftStore => false,
             Self::CondaPrefix
+            | Self::BaseCondaPrefix
             | Self::ProvidedPath
             | Self::ParentInterpreter
             | Self::ActiveEnvironment
@@ -1846,6 +1862,7 @@ impl VersionRequest {
             Self::Default => match source {
                 PythonSource::ParentInterpreter
                 | PythonSource::CondaPrefix
+                | PythonSource::BaseCondaPrefix
                 | PythonSource::ProvidedPath
                 | PythonSource::DiscoveredEnvironment
                 | PythonSource::ActiveEnvironment => Self::Any,
@@ -2256,7 +2273,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
@@ -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
@@ -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.
diff --git a/crates/uv-static/src/env_vars.rs b/crates/uv-static/src/env_vars.rs
@@ -337,6 +337,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";
+
     /// If set to `1` before a virtual environment is activated, then the
     /// virtual environment name will not be prepended to the terminal prompt.
     pub const VIRTUAL_ENV_DISABLE_PROMPT: &'static str = "VIRTUAL_ENV_DISABLE_PROMPT";
diff --git a/docs/configuration/environment.md b/docs/configuration/environment.md
@@ -139,6 +139,7 @@ uv respects the following environment variables:
   See [`PycInvalidationMode`](https://docs.python.org/3/library/py_compile.html#py_compile.PycInvalidationMode).
 - <a id="VIRTUAL_ENV"></a> [`VIRTUAL_ENV`](#VIRTUAL_ENV): Used to detect an activated virtual environment.
 - <a id="CONDA_PREFIX"></a> [`CONDA_PREFIX`](#CONDA_PREFIX): Used to detect an activated Conda environment.
+- <a id="CONDA_DEFAULT_ENV"></a> [`CONDA_DEFAULT_ENV`](#CONDA_DEFAULT_ENV): Used to determine if an active Conda environment is the base environment or not.
 - <a id="VIRTUAL_ENV_DISABLE_PROMPT"></a> [`VIRTUAL_ENV_DISABLE_PROMPT`](#VIRTUAL_ENV_DISABLE_PROMPT): If set to `1` before a virtual environment is activated, then the
   virtual environment name will not be prepended to the terminal prompt.
 - <a id="PROMPT"></a> [`PROMPT`](#PROMPT): Used to detect the use of the Windows Command Prompt (as opposed to PowerShell).
