Fix detection of Python on Codespaces

crates/pet-linux-global-python/src/lib.rs

@@ -1,23 +1,51 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use std::fs;
+use std::{
+    collections::HashMap,
+    fs,
+    path::{Path, PathBuf},
+    sync::{Arc, Mutex},
+    thread,
+};
 
-use log::warn;
 use pet_core::{
     python_environment::{PythonEnvironment, PythonEnvironmentBuilder, PythonEnvironmentCategory},
     reporter::Reporter,
     Locator,
 };
 use pet_fs::path::resolve_symlink;
-use pet_python_utils::{env::PythonEnv, executable::find_executables};
+use pet_python_utils::{
+    env::{PythonEnv, ResolvedPythonEnv},
+    executable::find_executables,
+};
 use pet_virtualenv::is_virtualenv;
 
-pub struct LinuxGlobalPython {}
+pub struct LinuxGlobalPython {
+    reported_executables: Arc<Mutex<HashMap<PathBuf, PythonEnvironment>>>,
+}
 
 impl LinuxGlobalPython {
     pub fn new() -> LinuxGlobalPython {
-        LinuxGlobalPython {}
+        LinuxGlobalPython {
+            reported_executables: Arc::new(
+                Mutex::new(HashMap::<PathBuf, PythonEnvironment>::new()),
+            ),
+        }
+    }
+
+    fn find_cached(&self, reporter: Option<&dyn Reporter>) {
+        if std::env::consts::OS == "macos" || std::env::consts::OS == "windows" {
+            return;
+        }
+        // Look through the /bin, /usr/bin, /usr/local/bin directories
+        thread::scope(|s| {
+            for bin in ["/bin", "/usr/bin", "/usr/local/bin"] {
+                s.spawn(move || {
+                    find_and_report_global_pythons_in(bin, reporter, &self.reported_executables);
+                });
+            }
+        });
     }
 }
 impl Default for LinuxGlobalPython {
@@ -47,41 +75,62 @@ impl Locator for LinuxGlobalPython {
         // If we do not have a version, then we cannot use this method.
         // Without version means we have not spawned the Python exe, thus do not have the real info.
         env.version.clone()?;
-        let prefix = env.prefix.clone()?;
         let executable = env.executable.clone();
 
-        // If prefix or version is not available, then we cannot use this method.
-        // 1. For files in /bin or /usr/bin, the prefix is always /usr
-        // 2. For files in /usr/local/bin, the prefix is always /usr/local
+        self.find_cached(None);
+
+        // We only support python environments in /bin, /usr/bin, /usr/local/bin
         if !executable.starts_with("/bin")
             && !executable.starts_with("/usr/bin")
             && !executable.starts_with("/usr/local/bin")
-            && !prefix.starts_with("/usr")
-            && !prefix.starts_with("/usr/local")
         {
             return None;
         }
 
-        // All known global linux are always installed in `/bin` or `/usr/bin` or `/usr/local/bin`
-        if executable.starts_with("/bin")
-            || executable.starts_with("/usr/bin")
-            || executable.starts_with("/usr/local/bin")
-        {
-            get_python_in_bin(env)
-        } else {
-            warn!(
-                    "Unknown Python exe ({:?}), not in any of the known locations /bin, /usr/bin, /usr/local/bin",
-                    executable
-                );
-            None
+        self.reported_executables
+            .lock()
+            .unwrap()
+            .get(&executable)
+            .cloned()
+    }
+
+    fn find(&self, reporter: &dyn Reporter) {
+        if std::env::consts::OS == "macos" || std::env::consts::OS == "windows" {
+            return;
         }
+        self.reported_executables.lock().unwrap().clear();
+        self.find_cached(Some(reporter))
     }
+}
 
-    fn find(&self, _reporter: &dyn Reporter) {
-        // No point looking in /usr/bin or /bin folder.
-        // We will end up looking in these global locations and spawning them in other parts.
-        // Here we cannot assume that anything in /usr/bin is a global Python, it could be a symlink or other.
-        // Safer approach is to just spawn it which we need to do to get the `sys.prefix`
+fn find_and_report_global_pythons_in(
+    bin: &str,
+    reporter: Option<&dyn Reporter>,
+    reported_executables: &Arc<Mutex<HashMap<PathBuf, PythonEnvironment>>>,
+) {
+    let python_executables = find_executables(Path::new(bin));
+
+    for exe in python_executables.clone().iter() {
+        if reported_executables.lock().unwrap().contains_key(exe) {
+            continue;
+        }
+        if let Some(resolved) = ResolvedPythonEnv::from(exe) {
+            if let Some(env) = get_python_in_bin(&resolved.to_python_env()) {
+                let mut reported_executables = reported_executables.lock().unwrap();
+                // env.symlinks = Some([symlinks, env.symlinks.clone().unwrap_or_default()].concat());
+                if let Some(symlinks) = &env.symlinks {
+                    for symlink in symlinks {
+                        reported_executables.insert(symlink.clone(), env.clone());
+                    }
+                }
+                if let Some(exe) = env.executable.clone() {
+                    reported_executables.insert(exe, env.clone());
+                }
+                if let Some(reporter) = reporter {
+                    reporter.report_environment(&env);
+                }
+            }
+        }
     }
 }
 
@@ -100,6 +149,7 @@ fn get_python_in_bin(env: &PythonEnv) -> Option<PythonEnvironment> {
     // Keep track of what the exe resolves to.
     // Will have a value only if the exe is in another dir
     // E.g. /bin/python3 might be a symlink to /usr/bin/python3.12
+    // Similarly /usr/local/python/current/bin/python might point to something like /usr/local/python/3.10.13/bin/python3.10
     // However due to legacy reasons we'll be treating these two as separate exes.
     // Hence they will be separate Python environments.
     let mut resolved_exe_is_from_another_dir = None;
@@ -119,6 +169,14 @@ fn get_python_in_bin(env: &PythonEnv) -> Option<PythonEnvironment> {
             resolved_exe_is_from_another_dir = Some(symlink);
         }
     }
+    if let Ok(symlink) = fs::canonicalize(&executable) {
+        // Ensure this is a symlink in the bin or usr/bin directory.
+        if symlink.starts_with(bin) {
+            symlinks.push(symlink);
+        } else {
+            resolved_exe_is_from_another_dir = Some(symlink);
+        }
+    }
 
     // Look for other symlinks in the same folder
     // We know that on linux there are sym links in the same folder as the exe.

crates/pet-poetry/src/config.rs

@@ -6,7 +6,7 @@ use std::{
     path::{Path, PathBuf},
 };
 
-use log::trace;
+use log::{error, trace};
 use pet_python_utils::platform_dirs::Platformdirs;
 
 use crate::env_variables::EnvVariables;
@@ -152,7 +152,7 @@ fn parse_contents(contents: &str) -> Option<ConfigToml> {
             })
         }
         Err(e) => {
-            eprintln!("Error parsing poetry toml file: {:?}", e);
+            error!("Error parsing poetry toml file: {:?}", e);
             None
         }
     }

crates/pet-poetry/src/pyproject_toml.rs

@@ -6,7 +6,7 @@ use std::{
     path::{Path, PathBuf},
 };
 
-use log::trace;
+use log::{error, trace};
 
 pub struct PyProjectToml {
     pub name: String,
@@ -41,7 +41,7 @@ fn parse_contents(contents: &str, file: &Path) -> Option<PyProjectToml> {
             name.map(|name| PyProjectToml::new(name, file.into()))
         }
         Err(e) => {
-            eprintln!("Error parsing toml file: {:?}", e);
+            error!("Error parsing toml file: {:?}", e);
             None
         }
     }

crates/pet/src/locators.rs

@@ -20,6 +20,7 @@ use pet_python_utils::env::{PythonEnv, ResolvedPythonEnv};
 use pet_venv::Venv;
 use pet_virtualenv::VirtualEnv;
 use pet_virtualenvwrapper::VirtualEnvWrapper;
+use std::path::PathBuf;
 use std::sync::Arc;
 
 pub fn create_locators(conda_locator: Arc<Conda>) -> Arc<Vec<Arc<dyn Locator>>> {
@@ -114,23 +115,72 @@ pub fn identify_python_environment_using_locators(
             // We have check all of the resolvers.
             // Telemetry point, failed to identify env here.
             warn!(
-                "Unknown Env ({:?}) in Path resolved as {:?} and reported as Unknown",
-                executable, resolved_env
+                "Unknown Env ({:?}) in Path resolved as {:?} and reported as {:?}",
+                executable,
+                resolved_env,
+                fallback_category.unwrap_or(PythonEnvironmentCategory::Unknown)
             );
-            let env = PythonEnvironmentBuilder::new(
-                fallback_category.unwrap_or(PythonEnvironmentCategory::Unknown),
-            )
-            .executable(Some(resolved_env.executable))
-            .prefix(Some(resolved_env.prefix))
-            .arch(Some(if resolved_env.is64_bit {
-                Architecture::X64
-            } else {
-                Architecture::X86
-            }))
-            .version(Some(resolved_env.version))
-            .build();
-            return Some(env);
+            return Some(create_unknown_env(resolved_env, fallback_category));
         }
     }
     None
 }
+
+fn create_unknown_env(
+    resolved_env: ResolvedPythonEnv,
+    fallback_category: Option<PythonEnvironmentCategory>,
+) -> PythonEnvironment {
+    // Find all the python exes in the same bin directory.
+
+    PythonEnvironmentBuilder::new(fallback_category.unwrap_or(PythonEnvironmentCategory::Unknown))
+        .symlinks(find_symlinks(&resolved_env.executable))
+        .executable(Some(resolved_env.executable))
+        .prefix(Some(resolved_env.prefix))
+        .arch(Some(if resolved_env.is64_bit {
+            Architecture::X64
+        } else {
+            Architecture::X86
+        }))
+        .version(Some(resolved_env.version))
+        .build()
+}
+
+#[cfg(unix)]
+fn find_symlinks(executable: &PathBuf) -> Option<Vec<PathBuf>> {
+    // Assume this is a python environment in /usr/bin/python.
+    // Now we know there can be other exes in the same directory as well, such as /usr/bin/python3.12 and that could be the same as /usr/bin/python
+    // However its possible /usr/bin/python is a symlink to /usr/local/bin/python3.12
+    // Either way, if both /usr/bin/python and /usr/bin/python3.12 point to the same exe (what ever it may be),
+    // then we know that both /usr/bin/python and /usr/bin/python3.12 are the same python environment.
+    // We use canonicalize to get the real path of the symlink.
+    // Only used in this case, see notes for resolve_symlink.
+
+    use pet_fs::path::resolve_symlink;
+    use pet_python_utils::executable::find_executables;
+    use std::fs;
+
+    let real_exe = resolve_symlink(executable).or(fs::canonicalize(executable).ok());
+
+    let bin = executable.parent()?;
+    // Make no assumptions that bin is always where exes are in linux
+    // No harm in supporting scripts as well.
+    if !bin.ends_with("bin") && !bin.ends_with("Scripts") && !bin.ends_with("scripts") {
+        return None;
+    }
+
+    let mut symlinks = vec![];
+    for exe in find_executables(bin) {
+        let symlink = resolve_symlink(&exe).or(fs::canonicalize(&exe).ok());
+        if symlink == real_exe {
+            symlinks.push(exe);
+        }
+    }
+    Some(symlinks)
+}
+
+#[cfg(windows)]
+fn find_symlinks(executable: &PathBuf) -> Option<Vec<PathBuf>> {
+    // In windows we will need to spawn the Python exe and then get the exes.
+    // Lets wait and see if this is necessary.
+    None
+}