Skip to content

Commit d24f9e7

Browse files
DonJayamanneDonJayamanne
authored
Fix winreg locator & add conda registry locator on windows (#23422)
1 parent 614aae6 commit d24f9e7

15 files changed

+139
-49
lines changed

‎native_locator/src/common_python.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ impl Locator for PythonOnPath<'_> {
3838
return None;
3939
}
4040
Some(PythonEnvironment {
41+
display_name: None,
4142
name: None,
4243
python_executable_path: Some(env.executable.clone()),
4344
version: env.version.clone(),

‎native_locator/src/conda.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -549,6 +549,7 @@ fn get_root_python_environment(path: &PathBuf, manager: &EnvManager) -> Option<P
549549
if let Some(package_info) = get_conda_package_json_path(&path, "python") {
550550
let conda_exe = manager.executable_path.to_str().unwrap().to_string();
551551
return Some(PythonEnvironment {
552+
display_name: None,
552553
name: None,
553554
category: messaging::PythonEnvironmentCategory::Conda,
554555
python_executable_path: Some(python_exe),
@@ -594,6 +595,7 @@ pub fn get_conda_environments_in_specified_path(
594595
envs.unwrap_or_default().iter().for_each(|env| {
595596
let exe = env.python_executable_path.clone();
596597
let env = PythonEnvironment::new(
598+
None,
597599
Some(env.name.clone()),
598600
exe.clone(),
599601
messaging::PythonEnvironmentCategory::Conda,
@@ -738,6 +740,7 @@ fn get_conda_environments_from_environments_txt_that_have_not_been_discovered(
738740
if let Some(env) = get_conda_environment_info(&PathBuf::from(env), false) {
739741
let exe = env.python_executable_path.clone();
740742
let env = PythonEnvironment::new(
743+
None,
741744
Some(env.name.clone()),
742745
exe.clone(),
743746
messaging::PythonEnvironmentCategory::Conda,

‎native_locator/src/conda_old.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -415,6 +415,7 @@ impl Locator for Conda<'_> {
415415
for env in envs {
416416
let executable = find_python_binary_path(Path::new(&env.path));
417417
let env = messaging::PythonEnvironment::new(
418+
None,
418419
Some(env.name.to_string()),
419420
executable.clone(),
420421
messaging::PythonEnvironmentCategory::Conda,

‎native_locator/src/homebrew.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ impl Locator for Homebrew<'_> {
6161
}
6262
reported.insert(exe.to_string_lossy().to_string());
6363
let env = crate::messaging::PythonEnvironment::new(
64+
None,
6465
None,
6566
Some(exe.clone()),
6667
crate::messaging::PythonEnvironmentCategory::Homebrew,

‎native_locator/src/main.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,24 +41,24 @@ fn main() {
4141
let pipenv_locator = pipenv::PipEnv::new();
4242
let mut path_locator = common_python::PythonOnPath::with(&environment);
4343
let mut conda_locator = conda::Conda::with(&environment);
44-
let mut pyenv_locator = pyenv::PyEnv::with(&environment, &mut conda_locator);
4544

4645
#[cfg(unix)]
4746
let mut homebrew_locator = homebrew::Homebrew::with(&environment);
4847
#[cfg(windows)]
4948
let mut windows_store = windows_store::WindowsStore::with(&environment);
5049
#[cfg(windows)]
51-
let mut windows_registry = windows_registry::WindowsRegistry::new();
50+
let mut windows_registry = windows_registry::WindowsRegistry::with(&mut conda_locator);
5251

5352
// Step 1: These environments take precedence over all others.
5453
// As they are very specific and guaranteed to be specific type.
54+
#[cfg(windows)]
55+
find_environments(&mut windows_registry, &mut dispatcher);
56+
let mut pyenv_locator = pyenv::PyEnv::with(&environment, &mut conda_locator);
5557
find_environments(&mut pyenv_locator, &mut dispatcher);
5658
#[cfg(unix)]
5759
find_environments(&mut homebrew_locator, &mut dispatcher);
5860
find_environments(&mut conda_locator, &mut dispatcher);
5961
#[cfg(windows)]
60-
find_environments(&mut windows_registry, &mut dispatcher);
61-
#[cfg(windows)]
6262
find_environments(&mut windows_store, &mut dispatcher);
6363

6464
// Step 2: Search in some global locations.

‎native_locator/src/messaging.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ pub enum PythonEnvironmentCategory {
9494
#[serde(rename_all = "camelCase")]
9595
#[derive(Debug)]
9696
pub struct PythonEnvironment {
97+
pub display_name: Option<String>,
9798
pub name: Option<String>,
9899
pub python_executable_path: Option<PathBuf>,
99100
pub category: PythonEnvironmentCategory,
@@ -110,6 +111,7 @@ pub struct PythonEnvironment {
110111

111112
impl PythonEnvironment {
112113
pub fn new(
114+
display_name: Option<String>,
113115
name: Option<String>,
114116
python_executable_path: Option<PathBuf>,
115117
category: PythonEnvironmentCategory,
@@ -120,6 +122,7 @@ impl PythonEnvironment {
120122
python_run_command: Option<Vec<String>>,
121123
) -> Self {
122124
Self {
125+
display_name,
123126
name,
124127
python_executable_path,
125128
category,
@@ -140,6 +143,7 @@ impl PythonEnvironment {
140143
project_path: PathBuf,
141144
) -> Self {
142145
Self {
146+
display_name: None,
143147
name: None,
144148
python_executable_path: python_executable_path.clone(),
145149
category: PythonEnvironmentCategory::Pipenv,

‎native_locator/src/pyenv.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ fn get_pure_python_environment(
105105
) -> Option<PythonEnvironment> {
106106
let version = get_pyenv_version(&path.file_name().unwrap().to_string_lossy().to_string())?;
107107
Some(messaging::PythonEnvironment::new(
108+
None,
108109
None,
109110
Some(executable.clone()),
110111
messaging::PythonEnvironmentCategory::Pyenv,
@@ -138,6 +139,7 @@ fn get_virtual_env_environment(
138139
let pyenv_cfg = find_and_parse_pyvenv_cfg(executable)?;
139140
let folder_name = path.file_name().unwrap().to_string_lossy().to_string();
140141
Some(messaging::PythonEnvironment::new(
142+
None,
141143
Some(folder_name),
142144
Some(executable.clone()),
143145
messaging::PythonEnvironmentCategory::PyenvVirtualEnv,

‎native_locator/src/venv.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ impl Locator for Venv {
2626
fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
2727
if is_venv(&env) {
2828
return Some(PythonEnvironment {
29+
display_name: None,
2930
name: Some(
3031
env.path
3132
.clone()

‎native_locator/src/virtualenv.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ impl Locator for VirtualEnv {
5757
fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
5858
if is_virtualenv(env) {
5959
return Some(PythonEnvironment {
60+
display_name: None,
6061
name: Some(
6162
env.path
6263
.clone()

‎native_locator/src/virtualenvwrapper.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ impl Locator for VirtualEnvWrapper<'_> {
8080
fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
8181
if is_virtualenvwrapper(env, self.environment) {
8282
return Some(PythonEnvironment {
83+
display_name: None,
8384
name: Some(
8485
env.path
8586
.clone()

‎native_locator/src/windows_registry.rs

Lines changed: 98 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,20 @@
11
// Copyright (c) Microsoft Corporation. All rights reserved.
22
// Licensed under the MIT License.
33

4+
#[cfg(windows)]
5+
use crate::conda::CondaLocator;
46
#[cfg(windows)]
57
use crate::locator::{Locator, LocatorResult};
68
#[cfg(windows)]
9+
use crate::messaging::EnvManager;
10+
#[cfg(windows)]
711
use crate::messaging::{PythonEnvironment, PythonEnvironmentCategory};
812
#[cfg(windows)]
913
use crate::utils::PythonEnv;
1014
#[cfg(windows)]
11-
use winreg::RegKey;
12-
#[cfg(windows)]
1315
use std::path::PathBuf;
16+
#[cfg(windows)]
17+
use winreg::RegKey;
1418

1519
#[cfg(windows)]
1620
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
1923

2024
let mut pythons = vec![];
2125
for key in company_key.enum_keys().filter_map(Result::ok) {
22-
let version_key = company_key.open_subkey(key).ok()?;
23-
let install_path_key = version_key.open_subkey("InstallPath").ok()?;
24-
let executable: String = install_path_key.get_value("ExecutablePath").ok()?;
25-
let version = version_key.get_value("Version").ok()?;
26-
27-
let env = PythonEnvironment::new(
28-
None,
29-
Some(PathBuf::from(executable.clone())),
30-
PythonEnvironmentCategory::WindowsRegistry,
31-
Some(version),
32-
None,
33-
None,
34-
None,
35-
Some(vec![executable.clone()]),
36-
);
37-
38-
pythons.push(env);
26+
if let Some(key) = company_key.open_subkey(key).ok() {
27+
if let Some(install_path_key) = key.open_subkey("InstallPath").ok() {
28+
let env_path: String = install_path_key.get_value("").ok().unwrap_or_default();
29+
let env_path = PathBuf::from(env_path);
30+
let env_path = if env_path.exists() {
31+
Some(env_path)
32+
} else {
33+
None
34+
};
35+
let executable: String = install_path_key
36+
.get_value("ExecutablePath")
37+
.ok()
38+
.unwrap_or_default();
39+
if executable.len() == 0 {
40+
continue;
41+
}
42+
let executable = PathBuf::from(executable);
43+
if !executable.exists() {
44+
continue;
45+
}
46+
let version: String = key.get_value("Version").ok().unwrap_or_default();
47+
let display_name: String = key.get_value("DisplayName").ok().unwrap_or_default();
48+
49+
let env = PythonEnvironment::new(
50+
Some(display_name),
51+
None,
52+
Some(executable.clone()),
53+
PythonEnvironmentCategory::WindowsRegistry,
54+
if version.len() > 0 {
55+
Some(version)
56+
} else {
57+
None
58+
},
59+
env_path,
60+
None,
61+
None,
62+
Some(vec![executable.to_string_lossy().to_string()]),
63+
);
64+
pythons.push(env);
65+
}
66+
}
3967
}
4068

4169
Some(pythons)
@@ -53,39 +81,77 @@ pub fn get_registry_pythons(company: &str) -> Option<Vec<PythonEnvironment>> {
5381
if let Some(hkcu_pythons) = get_registry_pythons_from_key(&hkcu, company) {
5482
pythons.extend(hkcu_pythons);
5583
}
56-
5784
Some(pythons)
5885
}
5986

6087
#[cfg(windows)]
61-
pub struct WindowsRegistry {}
88+
pub fn get_registry_pythons_anaconda(conda_locator: &mut dyn CondaLocator) -> LocatorResult {
89+
let hklm = winreg::RegKey::predef(winreg::enums::HKEY_LOCAL_MACHINE);
90+
let hkcu = winreg::RegKey::predef(winreg::enums::HKEY_CURRENT_USER);
91+
92+
let mut pythons = vec![];
93+
if let Some(hklm_pythons) = get_registry_pythons_from_key(&hklm, "ContinuumAnalytics") {
94+
pythons.extend(hklm_pythons);
95+
}
96+
if let Some(hkcu_pythons) = get_registry_pythons_from_key(&hkcu, "ContinuumAnalytics") {
97+
pythons.extend(hkcu_pythons);
98+
}
99+
100+
let mut environments: Vec<PythonEnvironment> = vec![];
101+
let mut managers: Vec<EnvManager> = vec![];
102+
103+
for env in pythons.iter() {
104+
if let Some(env_path) = env.clone().env_path {
105+
if let Some(mut result) = conda_locator.find_in(&env_path) {
106+
environments.append(&mut result.environments);
107+
managers.append(&mut result.managers);
108+
}
109+
}
110+
}
111+
112+
LocatorResult {
113+
managers,
114+
environments,
115+
}
116+
}
62117

63118
#[cfg(windows)]
64-
impl WindowsRegistry {
119+
pub struct WindowsRegistry<'a> {
120+
pub conda_locator: &'a mut dyn CondaLocator,
121+
}
122+
123+
#[cfg(windows)]
124+
impl WindowsRegistry<'_> {
65125
#[allow(dead_code)]
66-
pub fn new() -> WindowsRegistry {
67-
WindowsRegistry {}
126+
pub fn with<'a>(conda_locator: &'a mut impl CondaLocator) -> WindowsRegistry<'a> {
127+
WindowsRegistry { conda_locator }
68128
}
69129
}
70130

71131
#[cfg(windows)]
72-
impl Locator for WindowsRegistry {
73-
fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
132+
impl Locator for WindowsRegistry<'_> {
133+
fn resolve(&self, _env: &PythonEnv) -> Option<PythonEnvironment> {
74134
None
75135
}
76136

77137
fn find(&mut self) -> Option<LocatorResult> {
78-
let environments = get_registry_pythons("PythonCore")?;
79-
if environments.is_empty() {
138+
let mut environments: Vec<PythonEnvironment> = vec![];
139+
let mut managers: Vec<EnvManager> = vec![];
140+
141+
let mut result = get_registry_pythons("PythonCore").unwrap_or_default();
142+
environments.append(&mut result);
143+
144+
let mut result = get_registry_pythons_anaconda(self.conda_locator) ;
145+
environments.append(&mut result.environments);
146+
managers.append(&mut result.managers);
147+
148+
if environments.is_empty() && managers.is_empty() {
80149
None
81150
} else {
82151
Some(LocatorResult {
83-
managers: vec![],
152+
managers,
84153
environments,
85154
})
86155
}
87156
}
88157
}
89-
90-
// PythonCore
91-
// ContinuumAnalytics

‎native_locator/src/windows_store.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ impl Locator for WindowsStore<'_> {
5555
fn resolve(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
5656
if is_windows_python_executable(&env.executable) {
5757
return Some(PythonEnvironment {
58+
display_name: None,
5859
name: None,
5960
python_executable_path: Some(env.executable.clone()),
6061
version: None,

‎native_locator/tests/common_python_test.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ fn find_python_in_path_this() {
3333
assert_eq!(environments.len(), 1);
3434

3535
let env = PythonEnvironment {
36+
display_name: None,
3637
env_manager: None,
3738
project_path: None,
3839
name: None,

‎native_locator/tests/conda_test.rs

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ fn finds_two_conda_envs_from_txt() {
119119
tool: EnvManagerType::Conda,
120120
};
121121
let expected_conda_1 = PythonEnvironment {
122+
display_name: None,
122123
name: Some("one".to_string()),
123124
project_path: None,
124125
python_executable_path: Some(conda_1_exe.clone()),
@@ -133,9 +134,10 @@ fn finds_two_conda_envs_from_txt() {
133134
"-n".to_string(),
134135
"one".to_string(),
135136
"python".to_string(),
136-
]),
137-
};
138-
let expected_conda_2 = PythonEnvironment {
137+
]),
138+
};
139+
let expected_conda_2 = PythonEnvironment {
140+
display_name: None,
139141
name: Some("two".to_string()),
140142
project_path: None,
141143
python_executable_path: Some(conda_2_exe.clone()),

0 commit comments

Comments
 (0)