Updates to python locator (#23446)

The search algorithm has been documented here
https://github.com/microsoft/vscode-python/wiki/Python-Environment-Search

native_locator/src/common_python.rs

@@ -47,6 +47,7 @@ impl Locator for PythonOnPath<'_> {
             env_manager: None,
             project_path: None,
             python_run_command: Some(vec![env.executable.to_str().unwrap().to_string()]),
+            arch: None,
         })
     }
 

native_locator/src/conda.rs

@@ -646,6 +646,7 @@ fn get_root_python_environment(path: &PathBuf, manager: &EnvManager) -> Option<P
                 "python".to_string(),
             ]),
             project_path: None,
+            arch: None,
         });
     }
     None

native_locator/src/homebrew.rs

@@ -8,10 +8,9 @@ use crate::{
     utils::PythonEnv,
 };
 use regex::Regex;
-use std::{collections::HashSet, fs::DirEntry, path::PathBuf};
+use std::{collections::HashSet, path::PathBuf};
 
-fn is_symlinked_python_executable(path: &DirEntry) -> Option<PathBuf> {
-    let path = path.path();
+fn is_symlinked_python_executable(path: &PathBuf) -> Option<PathBuf> {
     let name = path.file_name()?.to_string_lossy();
     if !name.starts_with("python") || name.ends_with("-config") || name.ends_with("-build") {
         return None;
@@ -23,6 +22,149 @@ fn is_symlinked_python_executable(path: &DirEntry) -> Option<PathBuf> {
     Some(std::fs::canonicalize(path).ok()?)
 }
 
+fn get_homebrew_prefix_env_var(environment: &dyn Environment) -> Option<PathBuf> {
+    if let Some(homebrew_prefix) = environment.get_env_var("HOMEBREW_PREFIX".to_string()) {
+        let homebrew_prefix_bin = PathBuf::from(homebrew_prefix).join("bin");
+        if homebrew_prefix_bin.exists() {
+            return Some(homebrew_prefix_bin);
+        }
+    }
+    None
+}
+
+fn get_homebrew_prefix_bin(environment: &dyn Environment) -> Option<PathBuf> {
+    if let Some(homebrew_prefix) = get_homebrew_prefix_env_var(environment) {
+        return Some(homebrew_prefix);
+    }
+
+    // Homebrew install folders documented here https://docs.brew.sh/Installation
+    // /opt/homebrew for Apple Silicon,
+    // /usr/local for macOS Intel
+    // /home/linuxbrew/.linuxbrew for Linux
+    [
+        "/home/linuxbrew/.linuxbrew/bin",
+        "/opt/homebrew/bin",
+        "/usr/local/bin",
+    ]
+    .iter()
+    .map(|p| PathBuf::from(p))
+    .find(|p| p.exists())
+}
+
+fn get_env_path(python_exe_from_bin_dir: &PathBuf, resolved_file: &PathBuf) -> Option<PathBuf> {
+    // If the fully resolved file path contains the words `/homebrew/` or `/linuxbrew/`
+    // Then we know this is definitely a home brew version of python.
+    // And in these cases we can compute the sysprefix.
+
+    let resolved_file = resolved_file.to_str()?;
+    // 1. MacOS Silicon
+    if python_exe_from_bin_dir
+        .to_string_lossy()
+        .to_lowercase()
+        .starts_with("/opt/homebrew/bin/python")
+    {
+        // Resolved exe is something like `/opt/homebrew/Cellar/[email protected]/3.12.3/Frameworks/Python.framework/Versions/3.12/bin/python3.12`
+        let reg_ex = Regex::new("/opt/homebrew/Cellar/python@((\\d+\\.?)*)/(\\d+\\.?)*/Frameworks/Python.framework/Versions/(\\d+\\.?)*/bin/python(\\d+\\.?)*").unwrap();
+        let captures = reg_ex.captures(&resolved_file)?;
+        let version = captures.get(1).map(|m| m.as_str()).unwrap_or_default();
+        // SysPrefix- /opt/homebrew/opt/[email protected]/Frameworks/Python.framework/Versions/3.12
+        let sys_prefix = PathBuf::from(format!(
+            "/opt/homebrew/opt/python@{}/Frameworks/Python.framework/Versions/{}",
+            version, version
+        ));
+
+        return if sys_prefix.exists() {
+            Some(sys_prefix)
+        } else {
+            None
+        };
+    }
+
+    // 2. Linux
+    if python_exe_from_bin_dir
+        .to_string_lossy()
+        .to_lowercase()
+        .starts_with("/usr/local/bin/python")
+    {
+        // Resolved exe is something like `/home/linuxbrew/.linuxbrew/Cellar/[email protected]/3.12.3/bin/python3.12`
+        let reg_ex = Regex::new("/home/linuxbrew/.linuxbrew/Cellar/python@(\\d+\\.?\\d+\\.?)/(\\d+\\.?\\d+\\.?\\d+\\.?)/bin/python.*").unwrap();
+        let captures = reg_ex.captures(&resolved_file)?;
+        let version = captures.get(1).map(|m| m.as_str()).unwrap_or_default();
+        let full_version = captures.get(2).map(|m| m.as_str()).unwrap_or_default();
+        // SysPrefix- /home/linuxbrew/.linuxbrew/Cellar/[email protected]/3.12.3
+        let sys_prefix = PathBuf::from(format!(
+            "/home/linuxbrew/.linuxbrew/Cellar/python@{}/{}",
+            version, full_version
+        ));
+
+        return if sys_prefix.exists() {
+            Some(sys_prefix)
+        } else {
+            None
+        };
+    }
+
+    // 3. MacOS Intel
+    if python_exe_from_bin_dir
+        .to_string_lossy()
+        .to_lowercase()
+        .starts_with("/usr/local/bin/python")
+    {
+        // Resolved exe is something like `/usr/local/Cellar/[email protected]/3.12.3/Frameworks/Python.framework/Versions/3.12/bin/python3.12`
+        let reg_ex = Regex::new("/usr/local/Cellar/python@(\\d+\\.?\\d+\\.?)/(\\d+\\.?\\d+\\.?\\d+\\.?)/Frameworks/Python.framework/Versions/(\\d+\\.?\\d+\\.?)/bin/python.*").unwrap();
+        let captures = reg_ex.captures(&resolved_file)?;
+        let version = captures.get(1).map(|m| m.as_str()).unwrap_or_default();
+        let full_version = captures.get(2).map(|m| m.as_str()).unwrap_or_default();
+        // SysPrefix- /usr/local/Cellar/[email protected]/3.8.19/Frameworks/Python.framework/Versions/3.8
+        let sys_prefix = PathBuf::from(format!(
+            "/usr/local/Cellar/python@{}/{}/Frameworks/Python.framework/Versions/{}",
+            version, full_version, version
+        ));
+
+        return if sys_prefix.exists() {
+            Some(sys_prefix)
+        } else {
+            None
+        };
+    }
+    None
+}
+
+fn get_python_info(
+    python_exe_from_bin_dir: &PathBuf,
+    reported: &mut HashSet<String>,
+    python_version_regex: &Regex,
+) -> Option<PythonEnvironment> {
+    // Possible we do not have python3.12 or the like in bin directory
+    // & we have only python3, in that case we should add python3 to the list
+    if let Some(resolved_exe) = is_symlinked_python_executable(python_exe_from_bin_dir) {
+        let user_friendly_exe = python_exe_from_bin_dir;
+        let python_version = resolved_exe.to_string_lossy().to_string();
+        let version = match python_version_regex.captures(&python_version) {
+            Some(captures) => match captures.get(1) {
+                Some(version) => Some(version.as_str().to_string()),
+                None => None,
+            },
+            None => None,
+        };
+        if reported.contains(&resolved_exe.to_string_lossy().to_string()) {
+            return None;
+        }
+        reported.insert(resolved_exe.to_string_lossy().to_string());
+        return Some(PythonEnvironment::new(
+            None,
+            None,
+            Some(user_friendly_exe.clone()),
+            crate::messaging::PythonEnvironmentCategory::Homebrew,
+            version,
+            get_env_path(python_exe_from_bin_dir, &resolved_exe),
+            None,
+            Some(vec![user_friendly_exe.to_string_lossy().to_string()]),
+        ));
+    }
+    None
+}
+
 pub struct Homebrew<'a> {
     pub environment: &'a dyn Environment,
 }
@@ -34,64 +176,100 @@ impl Homebrew<'_> {
 }
 
 impl Locator for Homebrew<'_> {
-    fn resolve(&self, _env: &PythonEnv) -> Option<PythonEnvironment> {
-        None
+    fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
+        let python_regex = Regex::new(r"/(\d+\.\d+\.\d+)/").unwrap();
+        let exe = env.executable.clone();
+        let exe_file_name = exe.file_name()?;
+        let mut reported: HashSet<String> = HashSet::new();
+        if exe.starts_with("/opt/homebrew/bin/python")
+            || exe.starts_with("/opt/homebrew/Cellar/python@")
+            || exe.starts_with("/opt/homebrew/opt/python@")
+            || exe.starts_with("/opt/homebrew/opt/python")
+            || exe.starts_with("/opt/homebrew/Frameworks/Python.framework/Versions/")
+        {
+            // Symlink  - /opt/homebrew/bin/python3.12
+            // Symlink  - /opt/homebrew/opt/python3/bin/python3.12
+            // Symlink  - /opt/homebrew/Cellar/[email protected]/3.12.3/bin/python3.12
+            // Symlink  - /opt/homebrew/opt/[email protected]/bin/python3.12
+            // Symlink  - /opt/homebrew/Cellar/[email protected]/3.12.3/Frameworks/Python.framework/Versions/3.12/bin/python3.12
+            // Symlink  - /opt/homebrew/Cellar/[email protected]/3.12.3/Frameworks/Python.framework/Versions/Current/bin/python3.12
+            // Symlink  - /opt/homebrew/Frameworks/Python.framework/Versions/3.12/bin/python3.12
+            // Symlink  - /opt/homebrew/Frameworks/Python.framework/Versions/Current/bin/python3.12
+            // Real exe - /opt/homebrew/Cellar/[email protected]/3.12.3/Frameworks/Python.framework/Versions/3.12/bin/python3.12
+            // SysPrefix- /opt/homebrew/opt/[email protected]/Frameworks/Python.framework/Versions/3.12
+            get_python_info(
+                &PathBuf::from("/opt/homebrew/bin").join(exe_file_name),
+                &mut reported,
+                &python_regex,
+            )
+        } else if exe.starts_with("/usr/local/bin/python")
+            || exe.starts_with("/usr/local/opt/python@")
+            || exe.starts_with("/usr/local/Cellar/python@")
+        {
+            // Symlink  - /usr/local/bin/python3.8
+            // Symlink  - /usr/local/opt/[email protected]/bin/python3.8
+            // Symlink  - /usr/local/Cellar/[email protected]/3.8.19/bin/python3.8
+            // Real exe - /usr/local/Cellar/[email protected]/3.8.19/Frameworks/Python.framework/Versions/3.8/bin/python3.8
+            // SysPrefix- /usr/local/Cellar/[email protected]/3.8.19/Frameworks/Python.framework/Versions/3.8
+            get_python_info(
+                &PathBuf::from("/usr/local/bin").join(exe_file_name),
+                &mut reported,
+                &python_regex,
+            )
+        } else if exe.starts_with("/usr/local/bin/python")
+            || exe.starts_with("/home/linuxbrew/.linuxbrew/bin/python")
+            || exe.starts_with("/home/linuxbrew/.linuxbrew/opt/python@")
+            || exe.starts_with("/home/linuxbrew/.linuxbrew/Cellar/python")
+        {
+            // Symlink  - /usr/local/bin/python3.12
+            // Symlink  - /home/linuxbrew/.linuxbrew/bin/python3.12
+            // Symlink  - /home/linuxbrew/.linuxbrew/opt/[email protected]/bin/python3.12
+            // Real exe - /home/linuxbrew/.linuxbrew/Cellar/[email protected]/3.12.3/bin/python3.12
+            // SysPrefix- /home/linuxbrew/.linuxbrew/Cellar/[email protected]/3.12.3
+
+            get_python_info(
+                &PathBuf::from("/usr/local/bin").join(exe_file_name),
+                &mut reported,
+                &python_regex,
+            )
+        } else {
+            None
+        }
     }
 
     fn find(&mut self) -> Option<LocatorResult> {
-        let homebrew_prefix = self
-            .environment
-            .get_env_var("HOMEBREW_PREFIX".to_string())?;
-        let homebrew_prefix_bin = PathBuf::from(homebrew_prefix).join("bin");
+        let homebrew_prefix_bin = get_homebrew_prefix_bin(self.environment)?;
         let mut reported: HashSet<String> = HashSet::new();
         let python_regex = Regex::new(r"/(\d+\.\d+\.\d+)/").unwrap();
         let mut environments: Vec<PythonEnvironment> = vec![];
-        for file in std::fs::read_dir(homebrew_prefix_bin)
+        for file in std::fs::read_dir(&homebrew_prefix_bin)
             .ok()?
             .filter_map(Result::ok)
         {
-            if let Some(exe) = is_symlinked_python_executable(&file) {
-                let python_version = exe.to_string_lossy().to_string();
-                let version = match python_regex.captures(&python_version) {
-                    Some(captures) => match captures.get(1) {
-                        Some(version) => Some(version.as_str().to_string()),
-                        None => None,
-                    },
-                    None => None,
-                };
-                if reported.contains(&exe.to_string_lossy().to_string()) {
+            // If this file name is `python3`, then ignore this for now.
+            // We would prefer to use `python3.x` instead of `python3`.
+            // That way its more consistent and future proof
+            if let Some(file_name) = file.file_name().to_str() {
+                if file_name.to_lowercase() == "python3" {
                     continue;
                 }
-                let env_path = match exe.parent() {
-                    Some(path) => {
-                        if let Some(name) = path.file_name() {
-                            if name.to_ascii_lowercase() == "bin"
-                                || name.to_ascii_lowercase() == "Scripts"
-                            {
-                                Some(path.parent()?.to_path_buf())
-                            } else {
-                                Some(path.to_path_buf())
-                            }
-                        } else {
-                            None
-                        }
-                    }
-                    None => continue,
-                };
-                reported.insert(exe.to_string_lossy().to_string());
-                let env = crate::messaging::PythonEnvironment::new(
-                    None,
-                    None,
-                    Some(exe.clone()),
-                    crate::messaging::PythonEnvironmentCategory::Homebrew,
-                    version,
-                    env_path,
-                    None,
-                    Some(vec![exe.to_string_lossy().to_string()]),
-                );
+            }
+
+            if let Some(env) = get_python_info(&file.path(), &mut reported, &python_regex) {
                 environments.push(env);
             }
         }
+
+        // Possible we do not have python3.12 or the like in bin directory
+        // & we have only python3, in that case we should add python3 to the list
+        if let Some(env) = get_python_info(
+            &homebrew_prefix_bin.join("python3"),
+            &mut reported,
+            &python_regex,
+        ) {
+            environments.push(env);
+        }
+
         if environments.is_empty() {
             None
         } else {

native_locator/src/main.rs

@@ -37,7 +37,7 @@ fn main() {
 
     let virtualenv_locator = virtualenv::VirtualEnv::new();
     let venv_locator = venv::Venv::new();
-    let virtualenvwrapper_locator = virtualenvwrapper::VirtualEnvWrapper::with(&environment);
+    let mut virtualenvwrapper = virtualenvwrapper::VirtualEnvWrapper::with(&environment);
     let pipenv_locator = pipenv::PipEnv::new();
     let mut path_locator = common_python::PythonOnPath::with(&environment);
     let mut conda_locator = conda::Conda::with(&environment);
@@ -54,21 +54,35 @@ fn main() {
     #[cfg(windows)]
     find_environments(&mut windows_registry, &mut dispatcher);
     let mut pyenv_locator = pyenv::PyEnv::with(&environment, &mut conda_locator);
+    find_environments(&mut virtualenvwrapper, &mut dispatcher);
     find_environments(&mut pyenv_locator, &mut dispatcher);
     #[cfg(unix)]
     find_environments(&mut homebrew_locator, &mut dispatcher);
     find_environments(&mut conda_locator, &mut dispatcher);
     #[cfg(windows)]
     find_environments(&mut windows_store, &mut dispatcher);
 
-    // Step 2: Search in some global locations.
+    // Step 2: Search in some global locations for virtual envs.
     for env in list_global_virtual_envs(&environment).iter() {
         if dispatcher.was_environment_reported(&env) {
             continue;
         }
 
-        let _ = resolve_environment(&pipenv_locator, env, &mut dispatcher)
-            || resolve_environment(&virtualenvwrapper_locator, env, &mut dispatcher)
+        // First must be homebrew, as it is the most specific and supports symlinks
+        #[cfg(unix)]
+        let homebrew_result = resolve_environment(&homebrew_locator, env, &mut dispatcher);
+        #[cfg(unix)]
+        if homebrew_result {
+            continue;
+        }
+
+        let _ = // Pipeenv before virtualenvwrapper as it is more specific.
+            // Because pipenv environments are also virtualenvwrapper environments.
+            resolve_environment(&pipenv_locator, env, &mut dispatcher)
+            // Before venv, as all venvs are also virtualenvwrapper environments.
+            || resolve_environment(&virtualenvwrapper, env, &mut dispatcher)
+            // Before virtualenv as this is more specific.
+            // All venvs are also virtualenvs environments.
             || resolve_environment(&venv_locator, env, &mut dispatcher)
             || resolve_environment(&virtualenv_locator, env, &mut dispatcher);
     }