Support searching PATH, workspaces and some fixes (#41)

Author: DonJayamanne

.github/workflows/pr-check.yml

@@ -133,7 +133,7 @@ jobs:
         if: startsWith( matrix.os, 'ubuntu') || startsWith( matrix.os, 'macos')
         run: |
           pyenv install --list
-          pyenv install 3.12.3 3.8.19
+          pyenv install 3.12.4 3.8.19
         shell: bash
 
         # pyenv-win install list has not updated for a while
@@ -161,15 +161,15 @@ jobs:
         run: cargo fetch
         shell: bash
 
-      - name: Run Tests
-        run: cargo test --frozen --all-features -- --nocapture
-        shell: bash
-
       - name: Find Environments
         if: matrix.run_cli == 'yes'
         run: cargo run --release --target ${{ matrix.target }}
         shell: bash
 
+      - name: Run Tests
+        run: cargo test --frozen --all-features -- --nocapture
+        shell: bash
+
   builds:
     name: Builds
     runs-on: ${{ matrix.os }}

.gitignore

@@ -10,5 +10,9 @@ target/
 *.pdb
 
 # Directory with generated environments (generally created on CI)
-tmp
-.DS_Store
+.DS_Store
+
+# Testing directories
+.venv/
+tmp/
+temp/

Cargo.lock

@@ -9,6 +9,7 @@ msvc_spectre_libs = { version = "0.1.1", features = ["error"] }
 [dependencies]
 pet-fs = { path = "../pet-fs" }
 pet-python-utils = { path = "../pet-python-utils" }
+pet-virtualenv = { path = "../pet-virtualenv" }
 serde = { version = "1.0.152", features = ["derive"] }
 serde_json = "1.0.93"
 lazy_static = "1.4.0"

crates/pet-conda/Cargo.toml

@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use log::{error, trace, warn};
+use pet_fs::path::resolve_symlink;
+use std::path::PathBuf;
+
+#[derive(Debug, serde::Deserialize)]
+pub struct CondaInfo {
+    pub executable: PathBuf,
+    pub envs: Vec<PathBuf>,
+    pub conda_prefix: PathBuf,
+    pub conda_version: String,
+    pub envs_dirs: Vec<PathBuf>,
+}
+
+#[derive(Debug, serde::Deserialize)]
+pub struct CondaInfoJson {
+    pub envs: Option<Vec<PathBuf>>,
+    pub conda_prefix: Option<PathBuf>,
+    pub conda_version: Option<String>,
+    pub envs_dirs: Option<Vec<PathBuf>>,
+}
+
+impl CondaInfo {
+    pub fn from(executable: Option<PathBuf>) -> Option<CondaInfo> {
+        // let using_default = executable.is_none() || executable == Some("conda".into());
+        // Possible we got a symlink to the conda exe, first try to resolve that.
+        let executable = if cfg!(windows) {
+            executable.clone().unwrap_or("conda".into())
+        } else {
+            let executable = executable.unwrap_or("conda".into());
+            resolve_symlink(&executable).unwrap_or(executable)
+        };
+
+        let result = std::process::Command::new(&executable)
+            .arg("info")
+            .arg("--json")
+            .output();
+        trace!("Executing Conda: {:?} info --json", executable);
+        match result {
+            Ok(output) => {
+                if output.status.success() {
+                    let output = String::from_utf8_lossy(&output.stdout).to_string();
+                    match serde_json::from_str::<CondaInfoJson>(output.trim()) {
+                        Ok(info) => {
+                            let info = CondaInfo {
+                                executable: executable.clone(),
+                                envs: info.envs.unwrap_or_default().drain(..).collect(),
+                                conda_prefix: info.conda_prefix.unwrap_or_default(),
+                                conda_version: info.conda_version.unwrap_or_default(),
+                                envs_dirs: info.envs_dirs.unwrap_or_default().drain(..).collect(),
+                            };
+                            Some(info)
+                        }
+                        Err(err) => {
+                            error!(
+                                "Conda Execution for {:?} produced an output {:?} that could not be parsed as JSON",
+                                executable, err,
+                            );
+                            None
+                        }
+                    }
+                } else {
+                    let stderr = String::from_utf8_lossy(&output.stderr).to_string();
+                    // No point logging the message if conda is not installed or a custom conda exe wasn't provided.
+                    if executable.to_string_lossy() != "conda" {
+                        warn!(
+                            "Failed to get conda info using  {:?} ({:?}) {}",
+                            executable,
+                            output.status.code().unwrap_or_default(),
+                            stderr
+                        );
+                    }
+                    None
+                }
+            }
+            Err(err) => {
+                // No point logging the message if conda is not installed or a custom conda exe wasn't provided.
+                if executable.to_string_lossy() != "conda" {
+                    warn!("Failed to execute conda info {:?}", err);
+                }
+                None
+            }
+        }
+    }
+}

crates/pet-conda/src/conda_info.rs

@@ -14,7 +14,10 @@ use std::{
     thread,
 };
 
-pub fn get_conda_environment_paths(env_vars: &EnvVariables) -> Vec<PathBuf> {
+pub fn get_conda_environment_paths(
+    env_vars: &EnvVariables,
+    additional_env_dirs: &Vec<PathBuf>,
+) -> Vec<PathBuf> {
     let mut env_paths = thread::scope(|s| {
         let mut envs = vec![];
         for thread in [
@@ -26,6 +29,7 @@ pub fn get_conda_environment_paths(env_vars: &EnvVariables) -> Vec<PathBuf> {
             }),
             s.spawn(|| get_conda_environment_paths_from_conda_rc(env_vars)),
             s.spawn(|| get_conda_environment_paths_from_known_paths(env_vars)),
+            s.spawn(|| get_conda_environment_paths_from_additional_paths(additional_env_dirs)),
             s.spawn(|| get_known_conda_install_locations(env_vars)),
         ] {
             if let Ok(mut env_paths) = thread.join() {
@@ -102,6 +106,26 @@ fn get_conda_environment_paths_from_known_paths(env_vars: &EnvVariables) -> Vec<
     env_paths
 }
 
+fn get_conda_environment_paths_from_additional_paths(
+    additional_env_dirs: &Vec<PathBuf>,
+) -> Vec<PathBuf> {
+    let mut env_paths: Vec<PathBuf> = vec![];
+    for path in additional_env_dirs {
+        if let Ok(entries) = fs::read_dir(path) {
+            for entry in entries.filter_map(Result::ok) {
+                let path = entry.path();
+                if let Ok(meta) = fs::metadata(&path) {
+                    if meta.is_dir() {
+                        env_paths.push(path);
+                    }
+                }
+            }
+        }
+    }
+    env_paths.append(&mut additional_env_dirs.clone());
+    env_paths
+}
+
 pub fn get_environments(conda_dir: &Path) -> Vec<PathBuf> {
     let mut envs: Vec<PathBuf> = vec![];
 

crates/pet-conda/src/environment_locations.rs

@@ -1,16 +1,20 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+use conda_info::CondaInfo;
 use env_variables::EnvVariables;
 use environment_locations::{get_conda_environment_paths, get_environments};
 use environments::{get_conda_environment_info, CondaEnvironment};
-use log::error;
+use log::{error, warn};
 use manager::CondaManager;
 use pet_core::{
-    os_environment::Environment, python_environment::PythonEnvironment, reporter::Reporter,
+    os_environment::Environment,
+    python_environment::{PythonEnvironment, PythonEnvironmentCategory},
+    reporter::Reporter,
     Locator, LocatorResult,
 };
 use pet_python_utils::env::PythonEnv;
+use pet_virtualenv::is_virtualenv;
 use std::{
     collections::HashMap,
     path::{Path, PathBuf},
@@ -19,6 +23,7 @@ use std::{
 };
 use utils::is_conda_install;
 
+mod conda_info;
 pub mod conda_rc;
 pub mod env_variables;
 pub mod environment_locations;
@@ -29,8 +34,11 @@ pub mod utils;
 
 pub trait CondaLocator: Send + Sync {
     fn find_in(&self, path: &Path) -> Option<LocatorResult>;
+    fn find_with_conda_executable(&self, conda_executable: Option<PathBuf>) -> Option<()>;
 }
 pub struct Conda {
+    /// Directories where conda environments are found (env_dirs returned from `conda info --json`)
+    pub env_dirs: Arc<Mutex<Vec<PathBuf>>>,
     pub environments: Arc<Mutex<HashMap<PathBuf, PythonEnvironment>>>,
     pub managers: Arc<Mutex<HashMap<PathBuf, CondaManager>>>,
     pub env_vars: EnvVariables,
@@ -39,6 +47,7 @@ pub struct Conda {
 impl Conda {
     pub fn from(env: &dyn Environment) -> Conda {
         Conda {
+            env_dirs: Arc::new(Mutex::new(vec![])),
             environments: Arc::new(Mutex::new(HashMap::new())),
             managers: Arc::new(Mutex::new(HashMap::new())),
             env_vars: EnvVariables::from(env),
@@ -47,6 +56,40 @@ impl Conda {
 }
 
 impl CondaLocator for Conda {
+    fn find_with_conda_executable(&self, conda_executable: Option<PathBuf>) -> Option<()> {
+        let info = CondaInfo::from(conda_executable)?;
+        // Have we seen these envs, if yes, then nothing to do
+        let environments = self.environments.lock().unwrap().clone();
+        let mut new_envs = info
+            .envs
+            .clone()
+            .into_iter()
+            .filter(|p| !environments.contains_key(p))
+            .collect::<Vec<PathBuf>>();
+        if new_envs.is_empty() {
+            return None;
+        }
+
+        // Oh oh, we have new envs, lets find them.
+        let manager = CondaManager::from_info(&info.executable, &info)?;
+        for path in new_envs.iter() {
+            let mgr = manager.clone();
+            if let Some(env) = get_conda_environment_info(path, &Some(mgr.clone())) {
+                warn!(
+                    "Found a conda env {:?} using the conda exe {:?}",
+                    env.prefix, info.executable
+                );
+            }
+        }
+
+        // Also keep track of these directories for next time
+        let mut env_dirs = self.env_dirs.lock().unwrap();
+        env_dirs.append(&mut new_envs);
+
+        // Send telemetry
+
+        Some(())
+    }
     fn find_in(&self, conda_dir: &Path) -> Option<LocatorResult> {
         if !is_conda_install(conda_dir) {
             return None;
@@ -104,7 +147,18 @@ impl Conda {
 }
 
 impl Locator for Conda {
+    fn supported_categories(&self) -> Vec<PythonEnvironmentCategory> {
+        vec![PythonEnvironmentCategory::Conda]
+    }
     fn from(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
+        // Assume we create a virtual env from a conda python install,
+        // Doing this is a big no no, but people do this.
+        // Then the exe in the virtual env bin will be a symlink to the homebrew python install.
+        // Hence the first part of the condition will be true, but the second part will be false.
+        if is_virtualenv(env) {
+            return None;
+        }
+
         if let Some(ref path) = env.prefix {
             let mut environments = self.environments.lock().unwrap();
 
@@ -151,9 +205,10 @@ impl Locator for Conda {
         drop(environments);
 
         let env_vars = self.env_vars.clone();
+        let additional_paths = self.env_dirs.lock().unwrap().clone();
         thread::scope(|s| {
             // 1. Get a list of all know conda environments file paths
-            let possible_conda_envs = get_conda_environment_paths(&env_vars);
+            let possible_conda_envs = get_conda_environment_paths(&env_vars, &additional_paths);
             for path in possible_conda_envs {
                 s.spawn(move || {
                     // 2. Get the details of the conda environment