Skip to content

Commit 5523adc

Browse files
authored
Optimize workspace serach for .venv and .conda (#81)
1 parent 0fdb4b9 commit 5523adc

File tree

7 files changed

+170
-103
lines changed

7 files changed

+170
-103
lines changed

crates/pet-core/src/python_environment.rs

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,10 @@ pub struct PythonEnvironment {
7171
// Some of the known symlinks for the environment.
7272
// E.g. in the case of Homebrew there are a number of symlinks that are created.
7373
pub symlinks: Option<Vec<PathBuf>>,
74+
/// The folder/path that was searched to find this environment.
75+
/// Generally applies to workspace folder, and means that the environment was found in this folder or is related to this folder.
76+
/// Similar in meaqning to `project` but is more of a search path.
77+
pub search_path: Option<PathBuf>,
7478
}
7579

7680
impl Ord for PythonEnvironment {
@@ -109,6 +113,7 @@ impl Default for PythonEnvironment {
109113
project: None,
110114
arch: None,
111115
symlinks: None,
116+
search_path: None,
112117
}
113118
}
114119
}
@@ -159,6 +164,9 @@ impl std::fmt::Display for PythonEnvironment {
159164
if let Some(project) = &self.project {
160165
writeln!(f, " Project : {}", project.to_str().unwrap()).unwrap_or_default();
161166
}
167+
if let Some(search_path) = &self.search_path {
168+
writeln!(f, " Search Path : {}", search_path.to_str().unwrap()).unwrap_or_default();
169+
}
162170
if let Some(arch) = &self.arch {
163171
writeln!(f, " Architecture: {}", arch).unwrap_or_default();
164172
}
@@ -208,6 +216,7 @@ pub struct PythonEnvironmentBuilder {
208216
project: Option<PathBuf>,
209217
arch: Option<Architecture>,
210218
symlinks: Option<Vec<PathBuf>>,
219+
search_path: Option<PathBuf>,
211220
}
212221

213222
impl PythonEnvironmentBuilder {
@@ -223,6 +232,7 @@ impl PythonEnvironmentBuilder {
223232
project: None,
224233
arch: None,
225234
symlinks: None,
235+
search_path: None,
226236
}
227237
}
228238

@@ -285,6 +295,11 @@ impl PythonEnvironmentBuilder {
285295
self
286296
}
287297

298+
pub fn search_path(mut self, search_path: Option<PathBuf>) -> Self {
299+
self.search_path = search_path;
300+
self
301+
}
302+
288303
fn update_symlinks_and_exe(&mut self, symlinks: Option<Vec<PathBuf>>) {
289304
let mut all = vec![];
290305
if let Some(ref exe) = self.executable {
@@ -337,6 +352,7 @@ impl PythonEnvironmentBuilder {
337352
manager: self.manager,
338353
project: self.project,
339354
arch: self.arch,
355+
search_path: self.search_path,
340356
symlinks,
341357
}
342358
}

crates/pet-poetry/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,8 +140,8 @@ impl Locator for Poetry {
140140
}
141141
fn configure(&self, config: &Configuration) {
142142
if let Some(search_paths) = &config.search_paths {
143+
self.project_dirs.lock().unwrap().clear();
143144
if !search_paths.is_empty() {
144-
self.project_dirs.lock().unwrap().clear();
145145
self.project_dirs
146146
.lock()
147147
.unwrap()

crates/pet/src/find.rs

Lines changed: 79 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ pub fn find_and_report_envs(
109109
locators,
110110
false,
111111
&global_env_search_paths,
112+
None,
112113
);
113114
summary.lock().unwrap().find_path_time = start.elapsed();
114115
});
@@ -138,6 +139,7 @@ pub fn find_and_report_envs(
138139
locators,
139140
false,
140141
&global_env_search_paths,
142+
None,
141143
);
142144
summary.lock().unwrap().find_global_virtual_envs_time = start.elapsed();
143145
});
@@ -161,8 +163,6 @@ pub fn find_and_report_envs(
161163
search_paths,
162164
reporter,
163165
locators,
164-
0,
165-
1,
166166
);
167167
summary.lock().unwrap().find_search_paths_time = start.elapsed();
168168
});
@@ -173,61 +173,50 @@ pub fn find_and_report_envs(
173173
}
174174

175175
fn find_python_environments_in_workspace_folders_recursive(
176-
paths: Vec<PathBuf>,
176+
workspace_folders: Vec<PathBuf>,
177177
reporter: &dyn Reporter,
178178
locators: &Arc<Vec<Arc<dyn Locator>>>,
179-
depth: u32,
180-
max_depth: u32,
181179
) {
182180
thread::scope(|s| {
183-
// Find in cwd
184-
let paths1 = paths.clone();
185181
s.spawn(|| {
186-
find_python_environments(paths1, reporter, locators, true, &[]);
187-
188-
if depth >= max_depth {
189-
return;
190-
}
191-
192182
let bin = if cfg!(windows) { "Scripts" } else { "bin" };
193-
// If the folder has a bin or scripts, then ignore it, its most likely an env.
194-
// I.e. no point looking for python environments in a Python environment.
195-
let paths = paths
196-
.into_iter()
197-
.filter(|p| !p.join(bin).exists())
198-
.collect::<Vec<PathBuf>>();
183+
for workspace_folder in workspace_folders {
184+
find_python_environments_in_paths_with_locators(
185+
vec![
186+
// Possible this is a virtual env
187+
workspace_folder.clone(),
188+
// Optimize for finding these first.
189+
workspace_folder.join(".venv"),
190+
// Optimize for finding these first.
191+
workspace_folder.join(".conda"),
192+
],
193+
locators,
194+
reporter,
195+
true,
196+
&[],
197+
Some(workspace_folder.clone()),
198+
);
199+
200+
if workspace_folder.join(bin).exists() {
201+
// If the folder has a bin or scripts, then ignore it, its most likely an env.
202+
// I.e. no point looking for python environments in a Python environment.
203+
continue;
204+
}
199205

200-
for path in paths {
201-
if let Ok(reader) = fs::read_dir(&path) {
202-
let reader = reader
206+
if let Ok(reader) = fs::read_dir(&workspace_folder) {
207+
for folder in reader
203208
.filter_map(Result::ok)
204209
.filter(|d| d.file_type().is_ok_and(|f| f.is_dir()))
205210
.map(|p| p.path())
206-
.filter(should_search_for_environments_in_path);
207-
208-
// Take a batch of 20 items at a time.
209-
let reader = reader.fold(vec![], |f, a| {
210-
let mut f = f;
211-
if f.is_empty() {
212-
f.push(vec![a]);
213-
return f;
214-
}
215-
let last_item = f.last_mut().unwrap();
216-
if last_item.is_empty() || last_item.len() < 20 {
217-
last_item.push(a);
218-
return f;
219-
}
220-
f.push(vec![a]);
221-
f
222-
});
223-
224-
for entry in reader {
225-
find_python_environments_in_workspace_folders_recursive(
226-
entry,
211+
.filter(should_search_for_environments_in_path)
212+
{
213+
find_python_environments(
214+
vec![folder],
227215
reporter,
228216
locators,
229-
depth + 1,
230-
max_depth,
217+
true,
218+
&[],
219+
Some(workspace_folder.clone()),
231220
);
232221
}
233222
}
@@ -242,22 +231,23 @@ fn find_python_environments(
242231
locators: &Arc<Vec<Arc<dyn Locator>>>,
243232
is_workspace_folder: bool,
244233
global_env_search_paths: &[PathBuf],
234+
search_path: Option<PathBuf>,
245235
) {
246236
if paths.is_empty() {
247237
return;
248238
}
249239
thread::scope(|s| {
250-
let chunks = if is_workspace_folder { paths.len() } else { 1 };
251-
for item in paths.chunks(chunks) {
252-
let lst = item.to_vec().clone();
240+
for item in paths {
253241
let locators = locators.clone();
242+
let search_path = search_path.clone();
254243
s.spawn(move || {
255244
find_python_environments_in_paths_with_locators(
256-
lst,
245+
vec![item],
257246
&locators,
258247
reporter,
259248
is_workspace_folder,
260249
global_env_search_paths,
250+
search_path,
261251
);
262252
});
263253
}
@@ -270,55 +260,63 @@ fn find_python_environments_in_paths_with_locators(
270260
reporter: &dyn Reporter,
271261
is_workspace_folder: bool,
272262
global_env_search_paths: &[PathBuf],
263+
search_path: Option<PathBuf>,
273264
) {
274-
let executables = if is_workspace_folder {
275-
// If we're in a workspace folder, then we only need to look for bin/python or bin/python.exe
276-
// As workspace folders generally have either virtual env or conda env or the like.
277-
// They will not have environments that will ONLY have a file like `bin/python3`.
278-
// I.e. bin/python will almost always exist.
279-
paths
280-
.iter()
265+
for path in paths {
266+
let executables = if is_workspace_folder {
267+
// If we're in a workspace folder, then we only need to look for bin/python or bin/python.exe
268+
// As workspace folders generally have either virtual env or conda env or the like.
269+
// They will not have environments that will ONLY have a file like `bin/python3`.
270+
// I.e. bin/python will almost always exist.
271+
281272
// Paths like /Library/Frameworks/Python.framework/Versions/3.10/bin can end up in the current PATH variable.
282273
// Hence do not just look for files in a bin directory of the path.
283-
.flat_map(|p| find_executable(p))
284-
.filter_map(Option::Some)
285-
.collect::<Vec<PathBuf>>()
286-
} else {
287-
paths
288-
.iter()
274+
if let Some(executable) = find_executable(&path) {
275+
vec![executable]
276+
} else {
277+
vec![]
278+
}
279+
} else {
289280
// Paths like /Library/Frameworks/Python.framework/Versions/3.10/bin can end up in the current PATH variable.
290281
// Hence do not just look for files in a bin directory of the path.
291-
.flat_map(find_executables)
292-
.filter(|p| {
293-
// Exclude python2 on macOS
294-
if std::env::consts::OS == "macos" {
295-
return p.to_str().unwrap_or_default() != "/usr/bin/python2";
296-
}
297-
true
298-
})
299-
.collect::<Vec<PathBuf>>()
300-
};
282+
find_executables(path)
283+
.into_iter()
284+
.filter(|p| {
285+
// Exclude python2 on macOS
286+
if std::env::consts::OS == "macos" {
287+
return p.to_str().unwrap_or_default() != "/usr/bin/python2";
288+
}
289+
true
290+
})
291+
.collect::<Vec<PathBuf>>()
292+
};
301293

302-
identify_python_executables_using_locators(
303-
executables,
304-
locators,
305-
reporter,
306-
global_env_search_paths,
307-
);
294+
identify_python_executables_using_locators(
295+
executables,
296+
locators,
297+
reporter,
298+
global_env_search_paths,
299+
search_path.clone(),
300+
);
301+
}
308302
}
309303

310304
fn identify_python_executables_using_locators(
311305
executables: Vec<PathBuf>,
312306
locators: &Arc<Vec<Arc<dyn Locator>>>,
313307
reporter: &dyn Reporter,
314308
global_env_search_paths: &[PathBuf],
309+
search_path: Option<PathBuf>,
315310
) {
316311
for exe in executables.into_iter() {
317312
let executable = exe.clone();
318313
let env = PythonEnv::new(exe.to_owned(), None, None);
319-
if let Some(env) =
320-
identify_python_environment_using_locators(&env, locators, global_env_search_paths)
321-
{
314+
if let Some(env) = identify_python_environment_using_locators(
315+
&env,
316+
locators,
317+
global_env_search_paths,
318+
search_path.clone(),
319+
) {
322320
reporter.report_environment(&env);
323321
if let Some(manager) = env.manager {
324322
reporter.report_manager(&manager);

crates/pet/src/jsonrpc.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,11 +149,15 @@ pub fn handle_resolve(context: Arc<Context>, id: u32, params: Value) {
149149
match serde_json::from_value::<ResolveOptions>(params.clone()) {
150150
Ok(request_options) => {
151151
let executable = request_options.executable.clone();
152+
let search_paths = context.configuration.read().unwrap().clone().search_paths;
153+
let search_paths = search_paths.unwrap_or_default();
152154
// Start in a new thread, we can have multiple resolve requests.
153155
thread::spawn(move || {
154156
let now = SystemTime::now();
155157
trace!("Resolving env {:?}", executable);
156-
if let Some(result) = resolve_environment(&executable, &context.locators) {
158+
if let Some(result) =
159+
resolve_environment(&executable, &context.locators, search_paths)
160+
{
157161
if let Some(resolved) = result.resolved {
158162
// Gather telemetry of this resolved env and see what we got wrong.
159163
let _ = report_inaccuracies_identified_after_resolving(

0 commit comments

Comments
 (0)