From 41fc378661cd1287abe00a72e6f3eedafec6280f Mon Sep 17 00:00:00 2001 From: Zsolt Dollenstein Date: Mon, 22 Sep 2025 10:45:31 +0100 Subject: [PATCH] Discover uv workspaces --- Cargo.lock | 120 +++++++-- crates/pet-core/src/lib.rs | 1 + crates/pet-core/src/python_environment.rs | 2 + crates/pet-core/src/pyvenv_cfg.rs | 11 +- crates/pet-uv/Cargo.toml | 12 + crates/pet-uv/src/lib.rs | 287 ++++++++++++++++++++++ crates/pet/Cargo.toml | 1 + crates/pet/src/locators.rs | 2 + 8 files changed, 416 insertions(+), 20 deletions(-) create mode 100644 crates/pet-uv/Cargo.toml create mode 100644 crates/pet-uv/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 3e0837f9..174f744c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "ahash" @@ -234,13 +234,19 @@ dependencies = [ "allocator-api2", ] +[[package]] +name = "hashbrown" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5419bdc4f6a9207fbeba6d11b604d481addf78ecd10c11ad51e76c2f6482748d" + [[package]] name = "hashlink" version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8094feaf31ff591f651a2664fb9cfd92bba7a60ce3197265e9482ebe753c8f7" dependencies = [ - "hashbrown", + "hashbrown 0.14.5", ] [[package]] @@ -263,12 +269,12 @@ checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" [[package]] name = "indexmap" -version = "2.2.6" +version = "2.11.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "168fb715dda47215e360912c096649d23d58bf392ac62f73919e831745e40f26" +checksum = "4b0f83760fb341a774ed326568e19f5a863af4a952def8c39f9ab92fd95b88e5" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.16.0", ] [[package]] @@ -360,6 +366,7 @@ dependencies = [ "pet-python-utils", "pet-reporter", "pet-telemetry", + "pet-uv", "pet-venv", "pet-virtualenv", "pet-virtualenvwrapper", @@ -554,7 +561,7 @@ dependencies = [ "serde", "serde_json", "sha2", - "toml", + "toml 0.8.14", ] [[package]] @@ -617,6 +624,17 @@ dependencies = [ "regex", ] +[[package]] +name = "pet-uv" +version = "0.1.0" +dependencies = [ + "log", + "pet-core", + "pet-python-utils", + "serde", + "toml 0.9.7", +] + [[package]] name = "pet-venv" version = "0.1.0" @@ -685,9 +703,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.86" +version = "1.0.101" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e719e8df665df0d1c8fbfd238015744736151d4445ec0836b8e628aae103b77" +checksum = "89ae43fd86e4158d6db51ad8e2b80f313af9cc74f5c0e03ccb87de09998732de" dependencies = [ "unicode-ident", ] @@ -738,18 +756,28 @@ checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" [[package]] name = "serde" -version = "1.0.203" +version = "1.0.226" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7253ab4de971e72fb7be983802300c30b5a7f0c2e56fab8abfc6a214307c0094" +checksum = "0dca6411025b24b60bfa7ec1fe1f8e710ac09782dca409ee8237ba74b51295fd" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.226" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba2ba63999edb9dac981fb34b3e5c0d111a69b0924e253ed29d83f7c99e966a4" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.203" +version = "1.0.226" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "500cbc0ebeb6f46627f50f3f5811ccf6bf00643be300b4c3eabc0ef55dc5b5ba" +checksum = "8db53ae22f34573731bafa1db20f04027b2d25e02d8205921b569171699cdb33" dependencies = [ "proc-macro2", "quote", @@ -776,6 +804,15 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_spanned" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5417783452c2be558477e104686f7de5dae53dba813c28435e0e70f82d9b04ee" +dependencies = [ + "serde_core", +] + [[package]] name = "sha2" version = "0.10.8" @@ -795,9 +832,9 @@ checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" [[package]] name = "syn" -version = "2.0.67" +version = "2.0.106" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff8655ed1d86f3af4ee3fd3263786bc14245ad17c4c7e85ba7187fb3ae028c90" +checksum = "ede7c438028d4436d71104916910f5bb611972c5cfd7f89b8300a8186e6fada6" dependencies = [ "proc-macro2", "quote", @@ -820,11 +857,26 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f49eb2ab21d2f26bd6db7bf383edc527a7ebaee412d17af4d40fdccd442f335" dependencies = [ "serde", - "serde_spanned", - "toml_datetime", + "serde_spanned 0.6.6", + "toml_datetime 0.6.6", "toml_edit", ] +[[package]] +name = "toml" +version = "0.9.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00e5e5d9bf2475ac9d4f0d9edab68cc573dc2fd644b0dba36b0c30a92dd9eaa0" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned 1.0.2", + "toml_datetime 0.7.2", + "toml_parser", + "toml_writer", + "winnow 0.7.13", +] + [[package]] name = "toml_datetime" version = "0.6.6" @@ -834,6 +886,15 @@ dependencies = [ "serde", ] +[[package]] +name = "toml_datetime" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32f1085dec27c2b6632b04c80b3bb1b4300d6495d1e129693bdda7d91e72eec1" +dependencies = [ + "serde_core", +] + [[package]] name = "toml_edit" version = "0.22.14" @@ -842,11 +903,26 @@ checksum = "f21c7aaf97f1bd9ca9d4f9e73b0a6c74bd5afef56f2bc931943a6e1c37e04e38" dependencies = [ "indexmap", "serde", - "serde_spanned", - "toml_datetime", - "winnow", + "serde_spanned 0.6.6", + "toml_datetime 0.6.6", + "winnow 0.6.13", ] +[[package]] +name = "toml_parser" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cf893c33be71572e0e9aa6dd15e6677937abd686b066eac3f8cd3531688a627" +dependencies = [ + "winnow 0.7.13", +] + +[[package]] +name = "toml_writer" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d163a63c116ce562a22cda521fcc4d79152e7aba014456fb5eb442f6d6a10109" + [[package]] name = "typenum" version = "1.17.0" @@ -971,6 +1047,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "winnow" +version = "0.7.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21a0236b59786fed61e2a80582dd500fe61f18b5dca67a4a067d0bc9039339cf" + [[package]] name = "winreg" version = "0.55.0" diff --git a/crates/pet-core/src/lib.rs b/crates/pet-core/src/lib.rs index 8d123469..1768cae6 100644 --- a/crates/pet-core/src/lib.rs +++ b/crates/pet-core/src/lib.rs @@ -50,6 +50,7 @@ pub enum LocatorKind { Pixi, Poetry, PyEnv, + Uv, Venv, VirtualEnv, VirtualEnvWrapper, diff --git a/crates/pet-core/src/python_environment.rs b/crates/pet-core/src/python_environment.rs index a611ced8..361ec66b 100644 --- a/crates/pet-core/src/python_environment.rs +++ b/crates/pet-core/src/python_environment.rs @@ -23,6 +23,8 @@ pub enum PythonEnvironmentKind { MacCommandLineTools, LinuxGlobal, MacXCode, + Uv, + UvWorkspace, Venv, VirtualEnv, VirtualEnvWrapper, diff --git a/crates/pet-core/src/pyvenv_cfg.rs b/crates/pet-core/src/pyvenv_cfg.rs index 1dd454a5..28dbff54 100644 --- a/crates/pet-core/src/pyvenv_cfg.rs +++ b/crates/pet-core/src/pyvenv_cfg.rs @@ -23,6 +23,7 @@ pub struct PyVenvCfg { pub version_major: u64, pub version_minor: u64, pub prompt: Option, + pub file_path: PathBuf, } impl PyVenvCfg { @@ -31,12 +32,14 @@ impl PyVenvCfg { version_major: u64, version_minor: u64, prompt: Option, + file_path: PathBuf, ) -> Self { Self { version, version_major, version_minor, prompt, + file_path, } } pub fn find(path: &Path) -> Option { @@ -126,7 +129,13 @@ fn parse(file: &Path) -> Option { } match (version, version_major, version_minor) { - (Some(ver), Some(major), Some(minor)) => Some(PyVenvCfg::new(ver, major, minor, prompt)), + (Some(ver), Some(major), Some(minor)) => Some(PyVenvCfg::new( + ver, + major, + minor, + prompt, + file.to_path_buf(), + )), _ => None, } } diff --git a/crates/pet-uv/Cargo.toml b/crates/pet-uv/Cargo.toml new file mode 100644 index 00000000..b6e03729 --- /dev/null +++ b/crates/pet-uv/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "pet-uv" +version = "0.1.0" +edition = "2021" +license.workspace = true + +[dependencies] +pet-core = { path = "../pet-core" } +pet-python-utils = { path = "../pet-python-utils" } +serde = {version = "1.0.226", features = ["derive"]} +toml = "0.9.7" +log = "0.4.21" diff --git a/crates/pet-uv/src/lib.rs b/crates/pet-uv/src/lib.rs new file mode 100644 index 00000000..e658b33f --- /dev/null +++ b/crates/pet-uv/src/lib.rs @@ -0,0 +1,287 @@ +use std::{ + fs, + path::{Path, PathBuf}, + sync::{Arc, Mutex}, +}; + +use log::trace; +use pet_core::{ + env::PythonEnv, + python_environment::{PythonEnvironment, PythonEnvironmentBuilder, PythonEnvironmentKind}, + pyvenv_cfg::PyVenvCfg, + reporter::Reporter, + Configuration, Locator, LocatorKind, +}; +use pet_python_utils::executable::find_executables; +use serde::Deserialize; +pub struct Uv { + pub workspace_directories: Arc>>, +} + +/// Represents information stored in a `pyvenv.cfg` generated by uv +struct UvVenv { + uv_version: String, + python_version: String, + prompt: String, +} + +impl UvVenv { + fn maybe_from_file(file: &Path) -> Option { + let contents = fs::read_to_string(file).ok()?; + let mut uv_version = None; + let mut python_version = None; + let mut prompt = None; + for line in contents.lines() { + if let Some(uv_version_value) = line.trim_start().strip_prefix("uv = ") { + uv_version = Some(uv_version_value.trim_end().to_string()) + } + if let Some(version_info) = line.trim_start().strip_prefix("version_info = ") { + python_version = Some(version_info.to_string()); + } + if let Some(prompt_value) = line.trim_start().strip_prefix("prompt = ") { + prompt = Some(prompt_value.trim_end().to_string()); + } + } + Some(Self { + uv_version: uv_version?, + python_version: python_version?, + prompt: prompt?, + }) + } +} + +impl Uv { + pub fn new() -> Self { + Self { + workspace_directories: Arc::new(Mutex::new(Vec::new())), + } + } +} + +impl Locator for Uv { + fn get_kind(&self) -> LocatorKind { + LocatorKind::Uv + } + + fn supported_categories(&self) -> Vec { + vec![ + PythonEnvironmentKind::Uv, + PythonEnvironmentKind::UvWorkspace, + ] + } + + fn configure(&self, config: &Configuration) { + if let Some(workspace_directories) = config.workspace_directories.as_ref() { + let mut ws = self.workspace_directories.lock().unwrap(); + ws.clear(); + ws.extend(workspace_directories.iter().cloned()); + } + } + + fn try_from(&self, env: &PythonEnv) -> Option { + let cfg = env + .executable + .parent() + .and_then(|parent| PyVenvCfg::find(parent)) + .or_else(|| { + env.prefix + .as_ref() + .and_then(|prefix| PyVenvCfg::find(&prefix)) + })?; + let uv_venv = UvVenv::maybe_from_file(&cfg.file_path)?; + trace!( + "uv-managed venv found in {}, made by uv {}", + env.executable.display(), + uv_venv.uv_version + ); + let prefix = env.prefix.clone().or_else(|| { + env.executable + .parent() + .and_then(|p| p.parent().map(|pp| pp.to_path_buf())) + }); + let pyproject = prefix + .as_ref() + .and_then(|prefix| prefix.parent()) + .and_then(|parent| parse_pyproject_toml_in(parent)); + let kind = if pyproject + .and_then(|pyproject| pyproject.tool) + .and_then(|t| t.uv) + .and_then(|uv| uv.workspace) + .is_some() + { + PythonEnvironmentKind::UvWorkspace + } else { + PythonEnvironmentKind::Uv + }; + + Some( + PythonEnvironmentBuilder::new(Some(kind)) + .name(Some(uv_venv.prompt)) + .executable(Some(env.executable.clone())) + .version(Some(uv_venv.python_version)) + .symlinks(prefix.as_ref().map(|p| find_executables(p))) + .prefix(prefix) + .build(), + ) + } + + fn find(&self, reporter: &dyn Reporter) { + // look through workspace directories for uv-managed projects and any of their workspaces + let workspaces = self.workspace_directories.lock().unwrap().clone(); + for workspace in workspaces { + // TODO: maybe check for workspace in parent folders? + for env in list_envs_in_directory(&workspace) { + reporter.report_environment(&env); + } + } + } +} + +fn find_workspace(path: &Path) -> Option { + for candidate in path.ancestors() { + let pyproject = parse_pyproject_toml_in(&candidate); + if pyproject + .as_ref() + .and_then(|pp| pp.tool.as_ref()) + .and_then(|tool| tool.uv.as_ref()) + .and_then(|uv| uv.workspace.as_ref()) + .is_none() + { + continue; + } + trace!("Found workspace at {:?}", candidate); + let prefix = candidate.join(".venv"); + let pyvenv_cfg = prefix.join("pyvenv.cfg"); + if !pyvenv_cfg.exists() { + trace!( + "Workspace at {} does not have a virtual environment", + candidate.display() + ); + return None; + } + let unix_executable = prefix.join("bin/python"); + let windows_executable = prefix.join("Scripts/python.exe"); + let executable = if unix_executable.exists() { + Some(unix_executable) + } else if windows_executable.exists() { + Some(windows_executable) + } else { + None + }; + if let Some(uv_venv) = UvVenv::maybe_from_file(&pyvenv_cfg) { + return Some( + PythonEnvironmentBuilder::new(Some(PythonEnvironmentKind::UvWorkspace)) + .name(Some(uv_venv.prompt)) + .executable(executable) + .version(Some(uv_venv.python_version)) + .symlinks(Some(find_executables(&prefix))) + .prefix(Some(prefix)) + .build(), + ); + } else { + trace!( + "Workspace at {} does not have a uv-managed virtual environment", + candidate.display() + ); + } + return None; + } + None +} + +fn list_envs_in_directory(path: &Path) -> Vec { + let mut envs = Vec::new(); + let pyproject = parse_pyproject_toml_in(&path); + if pyproject.is_none() { + return envs; + } + let pyproject = pyproject.unwrap(); + let pyvenv_cfg = path.join(".venv/pyvenv.cfg"); + let prefix = path.join(".venv"); + let unix_executable = prefix.join("bin/python"); + let windows_executable = prefix.join("Scripts/python.exe"); + let executable = if unix_executable.exists() { + Some(unix_executable) + } else if windows_executable.exists() { + Some(windows_executable) + } else { + None + }; + if pyproject + .tool + .and_then(|t| t.uv) + .and_then(|uv| uv.workspace) + .is_some() + { + trace!("Workspace found in {}", path.display()); + let pyvenv_cfg = path.join(".venv/pyvenv.cfg"); + if let Some(uv_venv) = UvVenv::maybe_from_file(&pyvenv_cfg) { + trace!("uv-managed venv found for workspace in {}", path.display()); + let env = PythonEnvironmentBuilder::new(Some(PythonEnvironmentKind::UvWorkspace)) + .name(Some(uv_venv.prompt)) + .symlinks(Some(find_executables(&prefix))) + .prefix(Some(prefix)) + .executable(executable) + .version(Some(uv_venv.python_version)) + .build(); + envs.push(env); + } else { + trace!( + "No uv-managed venv found for workspace in {}", + path.display() + ); + } + // prioritize the workspace over the project if it's the same venv + } else if let Some(project) = pyproject.project { + if let Some(uv_venv) = UvVenv::maybe_from_file(&pyvenv_cfg) { + trace!("uv-managed venv found for project in {}", path.display()); + let env = PythonEnvironmentBuilder::new(Some(PythonEnvironmentKind::Uv)) + .name(Some(uv_venv.prompt)) + .symlinks(Some(find_executables(&prefix))) + .prefix(Some(prefix)) + .version(Some(uv_venv.python_version)) + .display_name(project.name) + .executable(executable) + .build(); + envs.push(env); + } else { + trace!("No uv-managed venv found in {}", path.display()); + } + if let Some(workspace) = path.parent().and_then(|p| find_workspace(p)) { + envs.push(workspace); + } + } + + envs +} + +fn parse_pyproject_toml_in(path: &Path) -> Option { + let contents = fs::read_to_string(path.join("pyproject.toml")).ok()?; + toml::from_str(&contents).ok() +} + +#[derive(Deserialize, Debug)] +struct PyProjectToml { + project: Option, + tool: Option, +} + +#[derive(Deserialize, Debug)] +struct Project { + name: Option, +} + +#[derive(Deserialize, Debug)] +struct Tool { + uv: Option, +} + +#[derive(Deserialize, Debug)] +struct ToolUv { + workspace: Option, +} + +#[cfg(test)] +mod tests { + use super::*; +} diff --git a/crates/pet/Cargo.toml b/crates/pet/Cargo.toml index ee2efdf7..99f7b23c 100644 --- a/crates/pet/Cargo.toml +++ b/crates/pet/Cargo.toml @@ -35,6 +35,7 @@ pet-virtualenv = { path = "../pet-virtualenv" } pet-pipenv = { path = "../pet-pipenv" } pet-telemetry = { path = "../pet-telemetry" } pet-global-virtualenvs = { path = "../pet-global-virtualenvs" } +pet-uv = { path = "../pet-uv" } log = "0.4.21" clap = { version = "4.5.4", features = ["derive", "cargo"] } serde = { version = "1.0.152", features = ["derive"] } diff --git a/crates/pet/src/locators.rs b/crates/pet/src/locators.rs index 5b6b3e81..c51302d2 100644 --- a/crates/pet/src/locators.rs +++ b/crates/pet/src/locators.rs @@ -19,6 +19,7 @@ use pet_pixi::Pixi; use pet_poetry::Poetry; use pet_pyenv::PyEnv; use pet_python_utils::env::ResolvedPythonEnv; +use pet_uv::Uv; use pet_venv::Venv; use pet_virtualenv::VirtualEnv; use pet_virtualenvwrapper::VirtualEnvWrapper; @@ -58,6 +59,7 @@ pub fn create_locators( // 6. Support for Virtual Envs // The order of these matter. // Basically PipEnv is a superset of VirtualEnvWrapper, which is a superset of Venv, which is a superset of VirtualEnv. + locators.push(Arc::new(Uv::new())); locators.push(poetry_locator); locators.push(Arc::new(PipEnv::from(environment))); locators.push(Arc::new(VirtualEnvWrapper::from(environment)));