Determine environments within the search instead of being asked to use a fallback (#76)

Author: DonJayamanne

crates/pet/src/find.rs

@@ -4,7 +4,6 @@
 use log::{info, trace, warn};
 use pet_conda::CondaLocator;
 use pet_core::os_environment::{Environment, EnvironmentApi};
-use pet_core::python_environment::PythonEnvironmentCategory;
 use pet_core::reporter::Reporter;
 use pet_core::{Configuration, Locator};
 use pet_env_var_path::get_search_paths_from_env_variables;
@@ -97,18 +96,19 @@ pub fn find_and_report_envs(
         s.spawn(|| {
             let start = std::time::Instant::now();
             let environment = EnvironmentApi::new();
-            let search_paths: Vec<PathBuf> = get_search_paths_from_env_variables(&environment);
+            let global_env_search_paths: Vec<PathBuf> =
+                get_search_paths_from_env_variables(&environment);
 
             trace!(
                 "Searching for environments in global folders: {:?}",
-                search_paths
+                global_env_search_paths
             );
             find_python_environments(
-                search_paths,
+                global_env_search_paths.clone(),
                 reporter,
                 locators,
                 false,
-                Some(PythonEnvironmentCategory::GlobalPaths),
+                &global_env_search_paths,
             );
             summary.lock().unwrap().find_path_time = start.elapsed();
         });
@@ -124,13 +124,21 @@ pub fn find_and_report_envs(
                 environment_paths,
             ]
             .concat();
+            let global_env_search_paths: Vec<PathBuf> =
+                get_search_paths_from_env_variables(&environment);
 
             trace!(
                 "Searching for environments in global venv folders: {:?}",
                 search_paths
             );
 
-            find_python_environments(search_paths, reporter, locators, false, None);
+            find_python_environments(
+                search_paths,
+                reporter,
+                locators,
+                false,
+                &global_env_search_paths,
+            );
             summary.lock().unwrap().find_global_virtual_envs_time = start.elapsed();
         });
         // Step 4: Find in workspace folders too.
@@ -175,7 +183,7 @@ fn find_python_environments_in_workspace_folders_recursive(
         // Find in cwd
         let paths1 = paths.clone();
         s.spawn(|| {
-            find_python_environments(paths1, reporter, locators, true, None);
+            find_python_environments(paths1, reporter, locators, true, &[]);
 
             if depth >= max_depth {
                 return;
@@ -233,7 +241,7 @@ fn find_python_environments(
     reporter: &dyn Reporter,
     locators: &Arc<Vec<Arc<dyn Locator>>>,
     is_workspace_folder: bool,
-    fallback_category: Option<PythonEnvironmentCategory>,
+    global_env_search_paths: &[PathBuf],
 ) {
     if paths.is_empty() {
         return;
@@ -249,7 +257,7 @@ fn find_python_environments(
                     &locators,
                     reporter,
                     is_workspace_folder,
-                    fallback_category,
+                    global_env_search_paths,
                 );
             });
         }
@@ -261,7 +269,7 @@ fn find_python_environments_in_paths_with_locators(
     locators: &Arc<Vec<Arc<dyn Locator>>>,
     reporter: &dyn Reporter,
     is_workspace_folder: bool,
-    fallback_category: Option<PythonEnvironmentCategory>,
+    global_env_search_paths: &[PathBuf],
 ) {
     let executables = if is_workspace_folder {
         // If we're in a workspace folder, then we only need to look for bin/python or bin/python.exe
@@ -291,20 +299,25 @@ fn find_python_environments_in_paths_with_locators(
             .collect::<Vec<PathBuf>>()
     };
 
-    identify_python_executables_using_locators(executables, locators, reporter, fallback_category);
+    identify_python_executables_using_locators(
+        executables,
+        locators,
+        reporter,
+        global_env_search_paths,
+    );
 }
 
 fn identify_python_executables_using_locators(
     executables: Vec<PathBuf>,
     locators: &Arc<Vec<Arc<dyn Locator>>>,
     reporter: &dyn Reporter,
-    fallback_category: Option<PythonEnvironmentCategory>,
+    global_env_search_paths: &[PathBuf],
 ) {
     for exe in executables.into_iter() {
         let executable = exe.clone();
         let env = PythonEnv::new(exe.to_owned(), None, None);
         if let Some(env) =
-            identify_python_environment_using_locators(&env, locators, fallback_category)
+            identify_python_environment_using_locators(&env, locators, global_env_search_paths)
         {
             reporter.report_environment(&env);
             if let Some(manager) = env.manager {

crates/pet/src/locators.rs

@@ -83,7 +83,7 @@ pub fn create_locators(conda_locator: Arc<Conda>) -> Arc<Vec<Arc<dyn Locator>>>
 pub fn identify_python_environment_using_locators(
     env: &PythonEnv,
     locators: &[Arc<dyn Locator>],
-    fallback_category: Option<PythonEnvironmentCategory>,
+    global_env_search_paths: &[PathBuf],
 ) -> Option<PythonEnvironment> {
     let executable = env.executable.clone();
     if let Some(env) = locators.iter().fold(
@@ -116,11 +116,25 @@ pub fn identify_python_environment_using_locators(
             // We have no idea what this is.
             // We have check all of the resolvers.
             // Telemetry point, failed to identify env here.
+            let mut fallback_category = None;
+
+            // If one of the symlinks are in the PATH variable, then we can treat this as a GlobalPath category.
+            let symlinks = [
+                resolved_env.symlink.clone().unwrap_or_default(),
+                vec![resolved_env.executable.clone(), executable.clone()],
+            ]
+            .concat();
+            for symlink in symlinks {
+                if let Some(bin) = symlink.parent() {
+                    if global_env_search_paths.contains(&bin.to_path_buf()) {
+                        fallback_category = Some(PythonEnvironmentCategory::GlobalPaths);
+                        break;
+                    }
+                }
+            }
             info!(
                 "Unknown Env ({:?}) in Path resolved as {:?} and reported as {:?}",
-                executable,
-                resolved_env,
-                fallback_category.unwrap_or(PythonEnvironmentCategory::Unknown)
+                executable, resolved_env, fallback_category
             );
             return Some(create_unknown_env(resolved_env, fallback_category));
         }

crates/pet/src/resolve.rs

@@ -3,12 +3,17 @@
 
 use std::{path::PathBuf, sync::Arc};
 
-use log::warn;
-use pet_core::{arch::Architecture, python_environment::PythonEnvironment, Locator};
+use log::{trace, warn};
+use pet_core::{
+    arch::Architecture, os_environment::EnvironmentApi, python_environment::PythonEnvironment,
+    Locator,
+};
+use pet_env_var_path::get_search_paths_from_env_variables;
 use pet_python_utils::env::{PythonEnv, ResolvedPythonEnv};
 
 use crate::locators::identify_python_environment_using_locators;
 
+#[derive(Debug)]
 pub struct ResolvedEnvironment {
     pub discovered: PythonEnvironment,
     pub resolved: Option<PythonEnvironment>,
@@ -20,17 +25,31 @@ pub fn resolve_environment(
 ) -> Option<ResolvedEnvironment> {
     // First check if this is a known environment
     let env = PythonEnv::new(executable.to_owned(), None, None);
-    if let Some(env) = identify_python_environment_using_locators(&env, locators, None) {
+    let environment = EnvironmentApi::new();
+    let global_env_search_paths: Vec<PathBuf> = get_search_paths_from_env_variables(&environment);
+
+    if let Some(env) =
+        identify_python_environment_using_locators(&env, locators, &global_env_search_paths)
+    {
         // Ok we got the environment.
         // Now try to resolve this fully, by spawning python.
         if let Some(ref executable) = env.executable {
             if let Some(info) = ResolvedPythonEnv::from(executable) {
+                trace!(
+                    "In resolve_environment, Resolved Python Exe {:?} as {:?}",
+                    executable,
+                    info
+                );
                 let discovered = env.clone();
                 let mut resolved = env.clone();
                 let mut symlinks = resolved.symlinks.clone().unwrap_or_default();
 
                 symlinks.push(info.executable.clone());
                 symlinks.append(&mut info.symlink.clone().unwrap_or_default());
+                symlinks.sort();
+                symlinks.dedup();
+
+                resolved.symlinks = Some(symlinks);
                 resolved.version = Some(info.version);
                 resolved.prefix = Some(info.prefix);
                 resolved.arch = Some(if info.is64_bit {

crates/pet/tests/ci_jupyter_container.rs

@@ -1,8 +1,21 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+use std::sync::Once;
+
 mod common;
 
+static INIT: Once = Once::new();
+
+/// Setup function that is only run once, even if called multiple times.
+fn setup() {
+    INIT.call_once(|| {
+        env_logger::builder()
+            .filter(None, log::LevelFilter::Trace)
+            .init();
+    });
+}
+
 #[cfg(unix)]
 #[cfg_attr(feature = "ci-jupyter-container", test)]
 #[allow(dead_code)]
@@ -19,6 +32,8 @@ fn verify_python_in_jupyter_contaner() {
     use pet_reporter::test;
     use std::{path::PathBuf, sync::Arc};
 
+    setup();
+
     let reporter = test::create_reporter();
     let environment = EnvironmentApi::new();
     let conda_locator = Arc::new(Conda::from(&environment));
@@ -121,23 +136,53 @@ fn verify_python_in_jupyter_contaner() {
             .iter()
             .find(|e| e.executable == env.executable)
             .expect(format!("Expected to find python environment {:?}", env.executable).as_str());
-        assert_eq!(python_env.executable, env.executable);
-        assert_eq!(python_env.category, env.category);
-        assert_eq!(python_env.symlinks, env.symlinks);
-        assert_eq!(python_env.manager, env.manager);
-        assert_eq!(python_env.name, env.name);
-        assert_eq!(python_env.version, env.version);
-        assert_eq!(python_env.arch, env.arch);
+        assert_eq!(
+            python_env.executable, env.executable,
+            "Expected exe to be same when comparing {:?} and {:?}",
+            python_env, env
+        );
+        assert_eq!(
+            python_env.category, env.category,
+            "Expected category to be same when comparing {:?} and {:?}",
+            python_env, env
+        );
+        assert_eq!(
+            python_env.symlinks, env.symlinks,
+            "Expected symlinks to be same when comparing {:?} and {:?}",
+            python_env, env
+        );
+        assert_eq!(
+            python_env.manager, env.manager,
+            "Expected manager to be same when comparing {:?} and {:?}",
+            python_env, env
+        );
+        assert_eq!(
+            python_env.name, env.name,
+            "Expected name to be same when comparing {:?} and {:?}",
+            python_env, env
+        );
+        assert_eq!(
+            python_env.version, env.version,
+            "Expected version to be same when comparing {:?} and {:?}",
+            python_env, env
+        );
+        assert_eq!(
+            python_env.arch, env.arch,
+            "Expected arch to be same when comparing {:?} and {:?}",
+            python_env, env
+        );
 
         // known issue https://github.com/microsoft/python-environment-tools/issues/64
         if env.executable == Some(PathBuf::from("/home/codespace/.python/current/bin/python")) {
             assert!(
                 python_env.prefix == Some(PathBuf::from("/home/codespace/.python/current"))
                     || python_env.prefix == Some(PathBuf::from("/usr/local/python/3.10.13")),
-                "Expected {:?} to be {:?} or {:?}",
+                "Expected {:?} to be {:?} or {:?} when comparing {:?} and {:?}",
                 python_env.prefix,
                 "/home/codespace/.python/current",
-                "/usr/local/python/3.10.13"
+                "/usr/local/python/3.10.13",
+                python_env,
+                env
             );
         }
     }