Fix winreg locator & add conda registry locator on windows (#23422)

native_locator/src/common_python.rs

@@ -38,6 +38,7 @@ impl Locator for PythonOnPath<'_> {
             return None;
         }
         Some(PythonEnvironment {
+            display_name: None,
             name: None,
             python_executable_path: Some(env.executable.clone()),
             version: env.version.clone(),

native_locator/src/conda.rs

@@ -549,6 +549,7 @@ fn get_root_python_environment(path: &PathBuf, manager: &EnvManager) -> Option<P
     if let Some(package_info) = get_conda_package_json_path(&path, "python") {
         let conda_exe = manager.executable_path.to_str().unwrap().to_string();
         return Some(PythonEnvironment {
+            display_name: None,
             name: None,
             category: messaging::PythonEnvironmentCategory::Conda,
             python_executable_path: Some(python_exe),
@@ -594,6 +595,7 @@ pub fn get_conda_environments_in_specified_path(
             envs.unwrap_or_default().iter().for_each(|env| {
                 let exe = env.python_executable_path.clone();
                 let env = PythonEnvironment::new(
+                    None,
                     Some(env.name.clone()),
                     exe.clone(),
                     messaging::PythonEnvironmentCategory::Conda,
@@ -738,6 +740,7 @@ fn get_conda_environments_from_environments_txt_that_have_not_been_discovered(
             if let Some(env) = get_conda_environment_info(&PathBuf::from(env), false) {
                 let exe = env.python_executable_path.clone();
                 let env = PythonEnvironment::new(
+                    None,
                     Some(env.name.clone()),
                     exe.clone(),
                     messaging::PythonEnvironmentCategory::Conda,

native_locator/src/conda_old.rs

@@ -415,6 +415,7 @@ impl Locator for Conda<'_> {
         for env in envs {
             let executable = find_python_binary_path(Path::new(&env.path));
             let env = messaging::PythonEnvironment::new(
+                None,
                 Some(env.name.to_string()),
                 executable.clone(),
                 messaging::PythonEnvironmentCategory::Conda,

native_locator/src/homebrew.rs

@@ -61,6 +61,7 @@ impl Locator for Homebrew<'_> {
                 }
                 reported.insert(exe.to_string_lossy().to_string());
                 let env = crate::messaging::PythonEnvironment::new(
+                    None,
                     None,
                     Some(exe.clone()),
                     crate::messaging::PythonEnvironmentCategory::Homebrew,

native_locator/src/main.rs

@@ -41,24 +41,24 @@ fn main() {
     let pipenv_locator = pipenv::PipEnv::new();
     let mut path_locator = common_python::PythonOnPath::with(&environment);
     let mut conda_locator = conda::Conda::with(&environment);
-    let mut pyenv_locator = pyenv::PyEnv::with(&environment, &mut conda_locator);
 
     #[cfg(unix)]
     let mut homebrew_locator = homebrew::Homebrew::with(&environment);
     #[cfg(windows)]
     let mut windows_store = windows_store::WindowsStore::with(&environment);
     #[cfg(windows)]
-    let mut windows_registry = windows_registry::WindowsRegistry::new();
+    let mut windows_registry = windows_registry::WindowsRegistry::with(&mut conda_locator);
 
     // Step 1: These environments take precedence over all others.
     // As they are very specific and guaranteed to be specific type.
+    #[cfg(windows)]
+    find_environments(&mut windows_registry, &mut dispatcher);
+    let mut pyenv_locator = pyenv::PyEnv::with(&environment, &mut conda_locator);
     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_registry, &mut dispatcher);
-    #[cfg(windows)]
     find_environments(&mut windows_store, &mut dispatcher);
 
     // Step 2: Search in some global locations.

native_locator/src/messaging.rs

@@ -94,6 +94,7 @@ pub enum PythonEnvironmentCategory {
 #[serde(rename_all = "camelCase")]
 #[derive(Debug)]
 pub struct PythonEnvironment {
+    pub display_name: Option<String>,
     pub name: Option<String>,
     pub python_executable_path: Option<PathBuf>,
     pub category: PythonEnvironmentCategory,
@@ -110,6 +111,7 @@ pub struct PythonEnvironment {
 
 impl PythonEnvironment {
     pub fn new(
+        display_name: Option<String>,
         name: Option<String>,
         python_executable_path: Option<PathBuf>,
         category: PythonEnvironmentCategory,
@@ -120,6 +122,7 @@ impl PythonEnvironment {
         python_run_command: Option<Vec<String>>,
     ) -> Self {
         Self {
+            display_name,
             name,
             python_executable_path,
             category,
@@ -140,6 +143,7 @@ impl PythonEnvironment {
         project_path: PathBuf,
     ) -> Self {
         Self {
+            display_name: None,
             name: None,
             python_executable_path: python_executable_path.clone(),
             category: PythonEnvironmentCategory::Pipenv,

native_locator/src/pyenv.rs

@@ -105,6 +105,7 @@ fn get_pure_python_environment(
 ) -> Option<PythonEnvironment> {
     let version = get_pyenv_version(&path.file_name().unwrap().to_string_lossy().to_string())?;
     Some(messaging::PythonEnvironment::new(
+        None,
         None,
         Some(executable.clone()),
         messaging::PythonEnvironmentCategory::Pyenv,
@@ -138,6 +139,7 @@ fn get_virtual_env_environment(
     let pyenv_cfg = find_and_parse_pyvenv_cfg(executable)?;
     let folder_name = path.file_name().unwrap().to_string_lossy().to_string();
     Some(messaging::PythonEnvironment::new(
+        None,
         Some(folder_name),
         Some(executable.clone()),
         messaging::PythonEnvironmentCategory::PyenvVirtualEnv,

native_locator/src/venv.rs

@@ -26,6 +26,7 @@ impl Locator for Venv {
     fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
         if is_venv(&env) {
             return Some(PythonEnvironment {
+                display_name: None,
                 name: Some(
                     env.path
                         .clone()

native_locator/src/virtualenv.rs

@@ -57,6 +57,7 @@ impl Locator for VirtualEnv {
     fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
         if is_virtualenv(env) {
             return Some(PythonEnvironment {
+                display_name: None,
                 name: Some(
                     env.path
                         .clone()

native_locator/src/virtualenvwrapper.rs

@@ -80,6 +80,7 @@ impl Locator for VirtualEnvWrapper<'_> {
     fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
         if is_virtualenvwrapper(env, self.environment) {
             return Some(PythonEnvironment {
+                display_name: None,
                 name: Some(
                     env.path
                         .clone()

native_locator/src/windows_registry.rs

@@ -1,16 +1,20 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#[cfg(windows)]
+use crate::conda::CondaLocator;
 #[cfg(windows)]
 use crate::locator::{Locator, LocatorResult};
 #[cfg(windows)]
+use crate::messaging::EnvManager;
+#[cfg(windows)]
 use crate::messaging::{PythonEnvironment, PythonEnvironmentCategory};
 #[cfg(windows)]
 use crate::utils::PythonEnv;
 #[cfg(windows)]
-use winreg::RegKey;
-#[cfg(windows)]
 use std::path::PathBuf;
+#[cfg(windows)]
+use winreg::RegKey;
 
 #[cfg(windows)]
 fn get_registry_pythons_from_key(hk: &RegKey, company: &str) -> Option<Vec<PythonEnvironment>> {
@@ -19,23 +23,47 @@ fn get_registry_pythons_from_key(hk: &RegKey, company: &str) -> Option<Vec<Pytho
 
     let mut pythons = vec![];
     for key in company_key.enum_keys().filter_map(Result::ok) {
-        let version_key = company_key.open_subkey(key).ok()?;
-        let install_path_key = version_key.open_subkey("InstallPath").ok()?;
-        let executable: String = install_path_key.get_value("ExecutablePath").ok()?;
-        let version = version_key.get_value("Version").ok()?;
-
-        let env = PythonEnvironment::new(
-            None,
-            Some(PathBuf::from(executable.clone())),
-            PythonEnvironmentCategory::WindowsRegistry,
-            Some(version),
-            None,
-            None,
-            None,
-            Some(vec![executable.clone()]),
-        );
-
-        pythons.push(env);
+        if let Some(key) = company_key.open_subkey(key).ok() {
+            if let Some(install_path_key) = key.open_subkey("InstallPath").ok() {
+                let env_path: String = install_path_key.get_value("").ok().unwrap_or_default();
+                let env_path = PathBuf::from(env_path);
+                let env_path = if env_path.exists() {
+                    Some(env_path)
+                } else {
+                    None
+                };
+                let executable: String = install_path_key
+                    .get_value("ExecutablePath")
+                    .ok()
+                    .unwrap_or_default();
+                if executable.len() == 0 {
+                    continue;
+                }
+                let executable = PathBuf::from(executable);
+                if !executable.exists() {
+                    continue;
+                }
+                let version: String = key.get_value("Version").ok().unwrap_or_default();
+                let display_name: String = key.get_value("DisplayName").ok().unwrap_or_default();
+
+                let env = PythonEnvironment::new(
+                    Some(display_name),
+                    None,
+                    Some(executable.clone()),
+                    PythonEnvironmentCategory::WindowsRegistry,
+                    if version.len() > 0 {
+                        Some(version)
+                    } else {
+                        None
+                    },
+                    env_path,
+                    None,
+                    None,
+                    Some(vec![executable.to_string_lossy().to_string()]),
+                );
+                pythons.push(env);
+            }
+        }
     }
 
     Some(pythons)
@@ -53,39 +81,77 @@ pub fn get_registry_pythons(company: &str) -> Option<Vec<PythonEnvironment>> {
     if let Some(hkcu_pythons) = get_registry_pythons_from_key(&hkcu, company) {
         pythons.extend(hkcu_pythons);
     }
-
     Some(pythons)
 }
 
 #[cfg(windows)]
-pub struct WindowsRegistry {}
+pub fn get_registry_pythons_anaconda(conda_locator: &mut dyn CondaLocator) -> LocatorResult {
+    let hklm = winreg::RegKey::predef(winreg::enums::HKEY_LOCAL_MACHINE);
+    let hkcu = winreg::RegKey::predef(winreg::enums::HKEY_CURRENT_USER);
+
+    let mut pythons = vec![];
+    if let Some(hklm_pythons) = get_registry_pythons_from_key(&hklm, "ContinuumAnalytics") {
+        pythons.extend(hklm_pythons);
+    }
+    if let Some(hkcu_pythons) = get_registry_pythons_from_key(&hkcu, "ContinuumAnalytics") {
+        pythons.extend(hkcu_pythons);
+    }
+
+    let mut environments: Vec<PythonEnvironment> = vec![];
+    let mut managers: Vec<EnvManager> = vec![];
+
+    for env in pythons.iter() {
+        if let Some(env_path) = env.clone().env_path {
+            if let Some(mut result) = conda_locator.find_in(&env_path) {
+                environments.append(&mut result.environments);
+                managers.append(&mut result.managers);
+            }
+        }
+    }
+
+    LocatorResult {
+        managers,
+        environments,
+    }
+}
 
 #[cfg(windows)]
-impl WindowsRegistry {
+pub struct WindowsRegistry<'a> {
+    pub conda_locator: &'a mut dyn CondaLocator,
+}
+
+#[cfg(windows)]
+impl WindowsRegistry<'_> {
     #[allow(dead_code)]
-    pub fn new() -> WindowsRegistry {
-        WindowsRegistry {}
+    pub fn with<'a>(conda_locator: &'a mut impl CondaLocator) -> WindowsRegistry<'a> {
+        WindowsRegistry { conda_locator }
     }
 }
 
 #[cfg(windows)]
-impl Locator for WindowsRegistry {
-    fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
+impl Locator for WindowsRegistry<'_> {
+    fn resolve(&self, _env: &PythonEnv) -> Option<PythonEnvironment> {
         None
     }
 
     fn find(&mut self) -> Option<LocatorResult> {
-        let environments = get_registry_pythons("PythonCore")?;
-        if environments.is_empty() {
+        let mut environments: Vec<PythonEnvironment> = vec![];
+        let mut managers: Vec<EnvManager> = vec![];
+
+        let mut result = get_registry_pythons("PythonCore").unwrap_or_default();
+        environments.append(&mut result);
+
+        let mut result = get_registry_pythons_anaconda(self.conda_locator) ;
+        environments.append(&mut result.environments);
+        managers.append(&mut result.managers);
+
+        if environments.is_empty() && managers.is_empty() {
             None
         } else {
             Some(LocatorResult {
-                managers: vec![],
+                managers,
                 environments,
             })
         }
     }
 }
-
-// PythonCore
-// ContinuumAnalytics

native_locator/src/windows_store.rs

@@ -55,6 +55,7 @@ impl Locator for WindowsStore<'_> {
     fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
         if is_windows_python_executable(&env.executable) {
             return Some(PythonEnvironment {
+                display_name: None,
                 name: None,
                 python_executable_path: Some(env.executable.clone()),
                 version: None,

native_locator/tests/common_python_test.rs

@@ -33,6 +33,7 @@ fn find_python_in_path_this() {
     assert_eq!(environments.len(), 1);
 
     let env = PythonEnvironment {
+        display_name: None,
         env_manager: None,
         project_path: None,
         name: None,

native_locator/tests/conda_test.rs

@@ -119,6 +119,7 @@ fn finds_two_conda_envs_from_txt() {
         tool: EnvManagerType::Conda,
     };
     let expected_conda_1 = PythonEnvironment {
+        display_name: None,
         name: Some("one".to_string()),
         project_path: None,
         python_executable_path: Some(conda_1_exe.clone()),
@@ -133,9 +134,10 @@ fn finds_two_conda_envs_from_txt() {
             "-n".to_string(),
             "one".to_string(),
             "python".to_string(),
-        ]),
-    };
-    let expected_conda_2 = PythonEnvironment {
+            ]),
+        };
+        let expected_conda_2 = PythonEnvironment {
+        display_name: None,
         name: Some("two".to_string()),
         project_path: None,
         python_executable_path: Some(conda_2_exe.clone()),