Skip to content

Commit f0c9407

Browse files
CopilotDonJayamanne
andcommitted
Implement UV virtual environment support
Co-authored-by: DonJayamanne <1948812+DonJayamanne@users.noreply.github.com>
1 parent a39a918 commit f0c9407

File tree

8 files changed

+226
-1
lines changed

8 files changed

+226
-1
lines changed

Cargo.lock

Lines changed: 12 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

crates/pet-uv/Cargo.toml

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
[package]
2+
name = "pet-uv"
3+
version = "0.1.0"
4+
edition = "2021"
5+
license = "MIT"
6+
7+
[dependencies]
8+
pet-core = { path = "../pet-core" }
9+
pet-fs = { path = "../pet-fs" }
10+
pet-conda = { path = "../pet-conda" }
11+
log = "0.4.21"

crates/pet-uv/src/lib.rs

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
// Copyright (c) Microsoft Corporation.
2+
// Licensed under the MIT License.
3+
4+
use pet_conda::utils::is_conda_env;
5+
use pet_fs::path::{expand_path, norm_case};
6+
use std::{fs, path::PathBuf};
7+
8+
/// Get the UV cache directory.
9+
/// UV uses the following priority order:
10+
/// 1. UV_CACHE_DIR environment variable
11+
/// 2. XDG cache directories on Unix / %LOCALAPPDATA% on Windows
12+
/// 3. Platform-specific cache directories
13+
fn get_uv_cache_dir(
14+
uv_cache_dir_env_var: Option<String>,
15+
xdg_cache_home: Option<String>,
16+
user_home: Option<PathBuf>,
17+
) -> Option<PathBuf> {
18+
// 1. Check UV_CACHE_DIR environment variable
19+
if let Some(cache_dir) = uv_cache_dir_env_var {
20+
let cache_dir = norm_case(expand_path(PathBuf::from(cache_dir)));
21+
if cache_dir.exists() {
22+
return Some(cache_dir);
23+
}
24+
}
25+
26+
// 2. Check XDG_CACHE_HOME on Unix
27+
if let Some(xdg_cache) = xdg_cache_home.map(|d| PathBuf::from(d).join("uv")) {
28+
if xdg_cache.exists() {
29+
return Some(xdg_cache);
30+
}
31+
}
32+
33+
// 3. Platform-specific cache directories
34+
if let Some(home) = user_home {
35+
let cache_dirs = if cfg!(target_os = "windows") {
36+
// On Windows: %LOCALAPPDATA%\uv
37+
vec![home.join("AppData").join("Local").join("uv")]
38+
} else if cfg!(target_os = "macos") {
39+
// On macOS: ~/Library/Caches/uv
40+
vec![home.join("Library").join("Caches").join("uv")]
41+
} else {
42+
// On other Unix systems: ~/.cache/uv
43+
vec![home.join(".cache").join("uv")]
44+
};
45+
46+
for cache_dir in cache_dirs {
47+
if cache_dir.exists() {
48+
return Some(cache_dir);
49+
}
50+
}
51+
}
52+
53+
None
54+
}
55+
56+
/// Get UV environment cache directories.
57+
/// UV stores virtual environments in {cache_dir}/environments-v2/
58+
fn get_uv_environment_dirs(
59+
uv_cache_dir_env_var: Option<String>,
60+
xdg_cache_home: Option<String>,
61+
user_home: Option<PathBuf>,
62+
) -> Vec<PathBuf> {
63+
let mut env_dirs = Vec::new();
64+
65+
if let Some(cache_dir) = get_uv_cache_dir(uv_cache_dir_env_var, xdg_cache_home, user_home) {
66+
let environments_dir = cache_dir.join("environments-v2");
67+
if environments_dir.exists() {
68+
env_dirs.push(environments_dir);
69+
}
70+
}
71+
72+
env_dirs
73+
}
74+
75+
/// List UV virtual environment paths.
76+
/// This function discovers UV cache directories and enumerates the virtual environments within them.
77+
/// It filters out conda environments to avoid conflicts.
78+
pub fn list_uv_virtual_envs_paths(
79+
uv_cache_dir_env_var: Option<String>,
80+
xdg_cache_home: Option<String>,
81+
user_home: Option<PathBuf>,
82+
) -> Vec<PathBuf> {
83+
let mut python_envs: Vec<PathBuf> = vec![];
84+
85+
for env_cache_dir in get_uv_environment_dirs(uv_cache_dir_env_var, xdg_cache_home, user_home) {
86+
if let Ok(dirs) = fs::read_dir(&env_cache_dir) {
87+
python_envs.append(
88+
&mut dirs
89+
.filter_map(Result::ok)
90+
.map(|e| e.path())
91+
.filter(|p| p.is_dir() && !is_conda_env(p))
92+
.collect(),
93+
);
94+
}
95+
}
96+
97+
python_envs.sort();
98+
python_envs.dedup();
99+
100+
python_envs
101+
}
102+
103+
#[cfg(test)]
104+
mod tests {
105+
use super::*;
106+
use std::fs;
107+
108+
#[test]
109+
fn test_uv_cache_dir_from_env_var() {
110+
let temp_dir = std::env::temp_dir().join("test_uv_cache");
111+
fs::create_dir_all(&temp_dir).unwrap();
112+
113+
let cache_dir = get_uv_cache_dir(
114+
Some(temp_dir.to_string_lossy().to_string()),
115+
None,
116+
None,
117+
);
118+
119+
assert_eq!(cache_dir, Some(temp_dir.clone()));
120+
fs::remove_dir_all(&temp_dir).ok();
121+
}
122+
123+
#[test]
124+
fn test_uv_environment_dirs() {
125+
let temp_dir = std::env::temp_dir().join("test_uv_env");
126+
let env_dir = temp_dir.join("environments-v2");
127+
fs::create_dir_all(&env_dir).unwrap();
128+
129+
let env_dirs = get_uv_environment_dirs(
130+
Some(temp_dir.to_string_lossy().to_string()),
131+
None,
132+
None,
133+
);
134+
135+
assert_eq!(env_dirs, vec![env_dir.clone()]);
136+
fs::remove_dir_all(&temp_dir).ok();
137+
}
138+
139+
#[test]
140+
fn test_list_uv_virtual_envs_paths() {
141+
let temp_dir = std::env::temp_dir().join("test_uv_list");
142+
let env_dir = temp_dir.join("environments-v2");
143+
let test_env = env_dir.join("test-venv");
144+
fs::create_dir_all(&test_env).unwrap();
145+
146+
let envs = list_uv_virtual_envs_paths(
147+
Some(temp_dir.to_string_lossy().to_string()),
148+
None,
149+
None,
150+
);
151+
152+
assert!(envs.contains(&test_env));
153+
fs::remove_dir_all(&temp_dir).ok();
154+
}
155+
}

crates/pet/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ pet-virtualenv = { path = "../pet-virtualenv" }
3535
pet-pipenv = { path = "../pet-pipenv" }
3636
pet-telemetry = { path = "../pet-telemetry" }
3737
pet-global-virtualenvs = { path = "../pet-global-virtualenvs" }
38+
pet-uv = { path = "../pet-uv" }
3839
log = "0.4.21"
3940
clap = { version = "4.5.4", features = ["derive", "cargo"] }
4041
serde = { version = "1.0.152", features = ["derive"] }

crates/pet/src/find.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use pet_core::{Configuration, Locator, LocatorKind};
1111
use pet_env_var_path::get_search_paths_from_env_variables;
1212
use pet_global_virtualenvs::list_global_virtual_envs_paths;
1313
use pet_pixi::is_pixi_env;
14+
use pet_uv::list_uv_virtual_envs_paths;
1415
use pet_python_utils::executable::{
1516
find_executable, find_executables, should_search_for_environments_in_path,
1617
};
@@ -165,6 +166,11 @@ pub fn find_and_report_envs(
165166
environment.get_env_var("XDG_DATA_HOME".into()),
166167
environment.get_user_home(),
167168
),
169+
list_uv_virtual_envs_paths(
170+
environment.get_env_var("UV_CACHE_DIR".into()),
171+
environment.get_env_var("XDG_CACHE_HOME".into()),
172+
environment.get_user_home(),
173+
),
168174
possible_environments,
169175
]
170176
.concat();

crates/pet/src/lib.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@ pub mod find;
2222
pub mod locators;
2323
pub mod resolve;
2424

25+
#[cfg(test)]
26+
mod tests;
27+
2528
#[derive(Debug, Clone)]
2629
pub struct FindOptions {
2730
pub print_list: bool,

crates/pet/src/tests.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
// Copyright (c) Microsoft Corporation.
2+
// Licensed under the MIT License.
3+
4+
mod test_uv_integration;
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
use std::fs;
2+
use pet_uv::list_uv_virtual_envs_paths;
3+
4+
#[test]
5+
fn test_uv_environment_discovery() {
6+
// Set up a temporary UV cache structure
7+
let temp_dir = std::env::temp_dir().join("test_pet_uv_integration");
8+
let cache_dir = temp_dir.join("uv");
9+
let env_dir = cache_dir.join("environments-v2");
10+
let test_env = env_dir.join("my-project-abc123-py3.12");
11+
let bin_dir = test_env.join("bin");
12+
13+
// Create the directory structure
14+
fs::create_dir_all(&bin_dir).unwrap();
15+
16+
// Create python executable and activate script to make it look like a virtual environment
17+
let python_exe = bin_dir.join("python");
18+
fs::write(&python_exe, "#!/bin/bash\necho 'python'").unwrap();
19+
let activate_script = bin_dir.join("activate");
20+
fs::write(&activate_script, "# Activate script").unwrap();
21+
22+
// Test UV path discovery
23+
let uv_paths = list_uv_virtual_envs_paths(
24+
Some(cache_dir.to_string_lossy().to_string()),
25+
None,
26+
None,
27+
);
28+
29+
// Verify that our test environment is discovered
30+
assert!(uv_paths.contains(&test_env), "UV environment should be discovered: {:?}", uv_paths);
31+
32+
// Clean up
33+
fs::remove_dir_all(&temp_dir).ok();
34+
}

0 commit comments

Comments
 (0)