Implement venv-uv environment type for UV-created virtual environments

Co-authored-by: karthiknadig <3840081+karthiknadig@users.noreply.github.com>

Author: Copilot

Cargo.lock

@@ -1,8 +1,8 @@
 ==> 2024-02-28 23:05:07 <==
-# cmd: /Users/donjayamanne/Development/vsc/python-environment-tools/crates/pet-conda/tests/unix/conda_env_without_manager_but_found_in_history/some_other_location/conda_install/bin/conda create -n conda1
+# cmd: /home/runner/work/python-environment-tools/python-environment-tools/crates/pet-conda/tests/unix/conda_env_without_manager_but_found_in_history/some_other_location/conda_install/bin/conda create -n conda1
 # conda version: 23.11.0
 ==> 2024-02-28 23:08:59 <==
-# cmd: /Users/donjayamanne/Development/vsc/python-environment-tools/crates/pet-conda/tests/unix/conda_env_without_manager_but_found_in_history/some_other_location/conda_install/bin/conda install -c conda-forge --name conda1 ipykernel -y
+# cmd: /home/runner/work/python-environment-tools/python-environment-tools/crates/pet-conda/tests/unix/conda_env_without_manager_but_found_in_history/some_other_location/conda_install/bin/conda install -c conda-forge --name conda1 ipykernel -y
 # conda version: 23.11.0
 +conda-forge/noarch::appnope-0.1.4-pyhd8ed1ab_0
 +conda-forge/noarch::asttokens-2.4.1-pyhd8ed1ab_0

crates/pet-conda/tests/unix/conda_env_without_manager_but_found_in_history/env_python_3/conda-meta/history

@@ -51,6 +51,7 @@ pub enum LocatorKind {
     Poetry,
     PyEnv,
     Venv,
+    VenvUv,
     VirtualEnv,
     VirtualEnvWrapper,
     WindowsRegistry,

crates/pet-core/src/lib.rs

@@ -24,6 +24,7 @@ pub enum PythonEnvironmentKind {
     LinuxGlobal,
     MacXCode,
     Venv,
+    VenvUv,
     VirtualEnv,
     VirtualEnvWrapper,
     WindowsStore,

crates/pet-core/src/python_environment.rs

@@ -23,6 +23,7 @@ pub struct PyVenvCfg {
     pub version_major: u64,
     pub version_minor: u64,
     pub prompt: Option<String>,
+    pub uv_version: Option<String>,
 }
 
 impl PyVenvCfg {
@@ -31,14 +32,19 @@ impl PyVenvCfg {
         version_major: u64,
         version_minor: u64,
         prompt: Option<String>,
+        uv_version: Option<String>,
     ) -> Self {
         Self {
             version,
             version_major,
             version_minor,
             prompt,
+            uv_version,
         }
     }
+    pub fn is_uv(&self) -> bool {
+        self.uv_version.is_some()
+    }
     pub fn find(path: &Path) -> Option<Self> {
         if let Some(ref file) = find(path) {
             parse(file)
@@ -99,6 +105,7 @@ fn parse(file: &Path) -> Option<PyVenvCfg> {
     let mut version_major: Option<u64> = None;
     let mut version_minor: Option<u64> = None;
     let mut prompt: Option<String> = None;
+    let mut uv_version: Option<String> = None;
 
     for line in contents.lines() {
         if version.is_none() {
@@ -120,13 +127,18 @@ fn parse(file: &Path) -> Option<PyVenvCfg> {
                 prompt = Some(p);
             }
         }
-        if version.is_some() && prompt.is_some() {
+        if uv_version.is_none() {
+            if let Some(uv_ver) = parse_uv_version(line) {
+                uv_version = Some(uv_ver);
+            }
+        }
+        if version.is_some() && prompt.is_some() && uv_version.is_some() {
             break;
         }
     }
 
     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, uv_version)),
         _ => None,
     }
 }
@@ -177,3 +189,72 @@ fn parse_prompt(line: &str) -> Option<String> {
     }
     None
 }
+
+fn parse_uv_version(line: &str) -> Option<String> {
+    let trimmed = line.trim();
+    if trimmed.starts_with("uv") {
+        if let Some(eq_idx) = trimmed.find('=') {
+            let mut version = trimmed[eq_idx + 1..].trim().to_string();
+            // Strip any leading or trailing single or double quotes
+            if version.starts_with('"') {
+                version = version.trim_start_matches('"').to_string();
+            }
+            if version.ends_with('"') {
+                version = version.trim_end_matches('"').to_string();
+            }
+            if version.starts_with('\'') {
+                version = version.trim_start_matches('\'').to_string();
+            }
+            if version.ends_with('\'') {
+                version = version.trim_end_matches('\'').to_string();
+            }
+            if !version.is_empty() {
+                return Some(version);
+            }
+        }
+    }
+    None
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::{path::PathBuf, fs};
+
+    #[test]
+    fn test_parse_uv_version() {
+        assert_eq!(parse_uv_version("uv = 0.8.14"), Some("0.8.14".to_string()));
+        assert_eq!(parse_uv_version("uv=0.8.14"), Some("0.8.14".to_string()));
+        assert_eq!(parse_uv_version("uv = \"0.8.14\""), Some("0.8.14".to_string()));
+        assert_eq!(parse_uv_version("uv = '0.8.14'"), Some("0.8.14".to_string()));
+        assert_eq!(parse_uv_version("version = 3.12.11"), None);
+        assert_eq!(parse_uv_version("prompt = test-env"), None);
+    }
+
+    #[test]
+    fn test_pyvenv_cfg_detects_uv() {
+        let temp_file = "/tmp/test_pyvenv_uv.cfg";
+        let contents = "home = /usr/bin/python3.12\nimplementation = CPython\nuv = 0.8.14\nversion_info = 3.12.11\ninclude-system-site-packages = false\nprompt = test-uv-env\n";
+        fs::write(temp_file, contents).unwrap();
+        
+        let cfg = parse(&PathBuf::from(temp_file)).unwrap();
+        assert!(cfg.is_uv());
+        assert_eq!(cfg.uv_version, Some("0.8.14".to_string()));
+        assert_eq!(cfg.prompt, Some("test-uv-env".to_string()));
+        
+        fs::remove_file(temp_file).ok();
+    }
+
+    #[test]
+    fn test_pyvenv_cfg_regular_venv() {
+        let temp_file = "/tmp/test_pyvenv_regular.cfg";
+        let contents = "home = /usr/bin/python3.12\ninclude-system-site-packages = false\nversion = 3.13.5\nexecutable = /usr/bin/python3.12\ncommand = python -m venv /path/to/env\n";
+        fs::write(temp_file, contents).unwrap();
+        
+        let cfg = parse(&PathBuf::from(temp_file)).unwrap();
+        assert!(!cfg.is_uv());
+        assert_eq!(cfg.uv_version, None);
+        
+        fs::remove_file(temp_file).ok();
+    }
+}

crates/pet-core/src/pyvenv_cfg.rs

@@ -0,0 +1,13 @@
+[package]
+name = "pet-venv-uv"
+version = "0.1.0"
+edition = "2021"
+license = "MIT"
+
+[target.'cfg(target_os = "windows")'.dependencies]
+msvc_spectre_libs = { version = "0.1.1", features = ["error"] }
+
+[dependencies]
+pet-core = { path = "../pet-core" }
+pet-python-utils = { path = "../pet-python-utils" }
+log = "0.4.21"

crates/pet-venv-uv/Cargo.toml

@@ -0,0 +1,130 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use std::path::Path;
+
+use pet_core::{
+    env::PythonEnv,
+    python_environment::{PythonEnvironment, PythonEnvironmentBuilder, PythonEnvironmentKind},
+    pyvenv_cfg::PyVenvCfg,
+    reporter::Reporter,
+    Locator, LocatorKind,
+};
+use pet_python_utils::executable::find_executables;
+use pet_python_utils::version;
+
+fn is_venv_uv_internal(env: &PythonEnv) -> Option<bool> {
+    // Check if there's a pyvenv.cfg file with uv entry
+    if let Some(cfg) = PyVenvCfg::find(env.executable.parent()?) {
+        return Some(cfg.is_uv());
+    }
+    if let Some(cfg) = PyVenvCfg::find(&env.prefix.clone()?) {
+        return Some(cfg.is_uv());
+    }
+    Some(false)
+}
+
+pub fn is_venv_uv(env: &PythonEnv) -> bool {
+    is_venv_uv_internal(env).unwrap_or_default()
+}
+
+pub fn is_venv_uv_dir(path: &Path) -> bool {
+    if let Some(cfg) = PyVenvCfg::find(path) {
+        cfg.is_uv()
+    } else {
+        false
+    }
+}
+
+pub struct VenvUv {}
+
+impl VenvUv {
+    pub fn new() -> VenvUv {
+        VenvUv {}
+    }
+}
+
+impl Default for VenvUv {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Locator for VenvUv {
+    fn get_kind(&self) -> LocatorKind {
+        LocatorKind::VenvUv
+    }
+
+    fn supported_categories(&self) -> Vec<PythonEnvironmentKind> {
+        vec![PythonEnvironmentKind::VenvUv]
+    }
+
+    fn try_from(&self, env: &PythonEnv) -> Option<PythonEnvironment> {
+        if is_venv_uv(env) {
+            let mut prefix = env.prefix.clone();
+            if prefix.is_none() {
+                prefix = Some(env.executable.parent()?.parent()?.to_path_buf());
+            }
+            let version = match env.version {
+                Some(ref v) => Some(v.clone()),
+                None => match &prefix {
+                    Some(prefix) => version::from_creator_for_virtual_env(prefix),
+                    None => None,
+                },
+            };
+            let mut symlinks = vec![];
+            if let Some(ref prefix) = prefix {
+                symlinks.append(&mut find_executables(prefix));
+            }
+
+            // Get the name from the prefix if it exists.
+            let cfg = PyVenvCfg::find(env.executable.parent()?)
+                .or_else(|| PyVenvCfg::find(&env.prefix.clone()?));
+            let name = cfg.and_then(|cfg| cfg.prompt);
+
+            Some(
+                PythonEnvironmentBuilder::new(Some(PythonEnvironmentKind::VenvUv))
+                    .name(name)
+                    .executable(Some(env.executable.clone()))
+                    .version(version)
+                    .prefix(prefix)
+                    .symlinks(Some(symlinks))
+                    .build(),
+            )
+        } else {
+            None
+        }
+    }
+
+    fn find(&self, _reporter: &dyn Reporter) {
+        // There are no common global locations for uv virtual environments.
+        // We expect the user of this class to call `is_compatible`
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::path::PathBuf;
+
+    #[test]
+    fn test_is_venv_uv_dir_detects_uv_environment() {
+        // This test checks if we can detect a UV environment from pyvenv.cfg
+        let uv_env_path = PathBuf::from("/tmp/test_uv_env");
+        assert!(is_venv_uv_dir(&uv_env_path), "Should detect UV environment");
+    }
+
+    #[test]
+    fn test_is_venv_uv_dir_does_not_detect_regular_environment() {
+        // This test checks if we can properly ignore regular venv environments
+        let regular_env_path = PathBuf::from("/tmp/test_regular_env");
+        assert!(!is_venv_uv_dir(&regular_env_path), "Should not detect regular venv as UV environment");
+    }
+
+    #[test]
+    fn test_is_venv_uv_dir_handles_nonexistent_environment() {
+        // This test checks if we handle non-existent environments gracefully
+        let nonexistent_path = PathBuf::from("/tmp/nonexistent_env");
+        assert!(!is_venv_uv_dir(&nonexistent_path), "Should not detect non-existent environment as UV");
+    }
+}

crates/pet-venv-uv/src/lib.rs

@@ -31,6 +31,7 @@ pet-linux-global-python = { path = "../pet-linux-global-python" }
 pet-mac-xcode = { path = "../pet-mac-xcode" }
 pet-mac-python-org = { path = "../pet-mac-python-org" }
 pet-venv = { path = "../pet-venv" }
+pet-venv-uv = { path = "../pet-venv-uv" }
 pet-virtualenv = { path = "../pet-virtualenv" }
 pet-pipenv = { path = "../pet-pipenv" }
 pet-telemetry = { path = "../pet-telemetry" }

crates/pet/Cargo.toml

@@ -20,6 +20,7 @@ use pet_poetry::Poetry;
 use pet_pyenv::PyEnv;
 use pet_python_utils::env::ResolvedPythonEnv;
 use pet_venv::Venv;
+use pet_venv_uv::VenvUv;
 use pet_virtualenv::VirtualEnv;
 use pet_virtualenvwrapper::VirtualEnvWrapper;
 use std::path::PathBuf;
@@ -57,10 +58,11 @@ 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.
+    // Basically PipEnv is a superset of VirtualEnvWrapper, which is a superset of VenvUv, which is a superset of Venv, which is a superset of VirtualEnv.
     locators.push(poetry_locator);
     locators.push(Arc::new(PipEnv::from(environment)));
     locators.push(Arc::new(VirtualEnvWrapper::from(environment)));
+    locators.push(Arc::new(VenvUv::new()));
     locators.push(Arc::new(Venv::new()));
     // VirtualEnv is the most generic, hence should be the last.
     locators.push(Arc::new(VirtualEnv::new()));