[ty] Support ephemeral uv virtual environments

Author: AlexWaygood

crates/ty_python_semantic/src/module_resolver/resolver.rs

@@ -4,6 +4,7 @@ use std::iter::FusedIterator;
 use std::str::Split;
 
 use compact_str::format_compact;
+use indexmap::IndexSet;
 use rustc_hash::{FxBuildHasher, FxHashSet};
 
 use ruff_db::files::{File, FilePath, FileRootKind};
@@ -289,7 +290,7 @@ impl SearchPaths {
                             virtual_env_path,
                             error
                         );
-                        vec![]
+                        IndexSet::default()
                     };
 
                     match PythonEnvironment::new(
@@ -304,7 +305,7 @@ impl SearchPaths {
                     }
                 } else {
                     tracing::debug!("No virtual environment found");
-                    vec![]
+                    IndexSet::default()
                 }
             }
 

crates/ty_python_semantic/src/site_packages.rs

@@ -14,6 +14,7 @@ use std::io;
 use std::num::NonZeroUsize;
 use std::ops::Deref;
 
+use indexmap::IndexSet;
 use ruff_db::system::{System, SystemPath, SystemPathBuf};
 use ruff_python_ast::PythonVersion;
 
@@ -35,7 +36,7 @@ impl PythonEnvironment {
 
         // Attempt to inspect as a virtual environment first
         // TODO(zanieb): Consider avoiding the clone here by checking for `pyvenv.cfg` ahead-of-time
-        match VirtualEnvironment::new(path.clone(), origin, system) {
+        match VirtualEnvironment::new(path.clone(), system) {
             Ok(venv) => Ok(Self::Virtual(venv)),
             // If there's not a `pyvenv.cfg` marker, attempt to inspect as a system environment
             //
@@ -54,7 +55,7 @@ impl PythonEnvironment {
     pub(crate) fn site_packages_directories(
         &self,
         system: &dyn System,
-    ) -> SitePackagesDiscoveryResult<Vec<SystemPathBuf>> {
+    ) -> SitePackagesDiscoveryResult<IndexSet<SystemPathBuf>> {
         match self {
             Self::Virtual(env) => env.site_packages_directories(system),
             Self::System(env) => env.site_packages_directories(system),
@@ -111,12 +112,19 @@ pub(crate) struct VirtualEnvironment {
     /// This field will be `None` if so.
     version: Option<PythonVersion>,
     implementation: PythonImplementation,
+
+    /// If this virtual environment was created using uv,
+    /// it may be an "ephemeral" virtual environment that dynamically adds the `site-packages`
+    /// directories of its parent environment to `sys.path` at runtime.
+    /// Newer versions of uv record the parent environment in the `pyvenv.cfg` file;
+    /// we'll want to add the `site-packages` directories of the parent environment
+    /// as search paths as well as the `site-packages` directories of this virtual environment.
+    parent_environment: Option<Box<VirtualEnvironment>>,
 }
 
 impl VirtualEnvironment {
     pub(crate) fn new(
         path: SysPrefixPath,
-        origin: SysPrefixPathOrigin,
         system: &dyn System,
     ) -> SitePackagesDiscoveryResult<Self> {
         fn pyvenv_cfg_line_number(index: usize) -> NonZeroUsize {
@@ -128,12 +136,14 @@ impl VirtualEnvironment {
 
         let pyvenv_cfg = system
             .read_to_string(&pyvenv_cfg_path)
-            .map_err(|io_err| SitePackagesDiscoveryError::NoPyvenvCfgFile(origin, io_err))?;
+            .map_err(|io_err| SitePackagesDiscoveryError::NoPyvenvCfgFile(path.origin, io_err))?;
 
         let mut include_system_site_packages = false;
         let mut base_executable_home_path = None;
         let mut version_info_string = None;
         let mut implementation = PythonImplementation::Unknown;
+        let mut created_with_uv = false;
+        let mut parent_environment = None;
 
         // A `pyvenv.cfg` file *looks* like a `.ini` file, but actually isn't valid `.ini` syntax!
         // The Python standard-library's `site` module parses these files by splitting each line on
@@ -178,6 +188,8 @@ impl VirtualEnvironment {
                             _ => PythonImplementation::Unknown,
                         };
                     }
+                    "uv" => created_with_uv = true,
+                    "extends-environment" => parent_environment = Some(value),
                     _ => continue,
                 }
             }
@@ -196,11 +208,35 @@ impl VirtualEnvironment {
         let base_executable_home_path = PythonHomePath::new(base_executable_home_path, system)
             .map_err(|io_err| {
                 SitePackagesDiscoveryError::PyvenvCfgParseError(
-                    pyvenv_cfg_path,
+                    pyvenv_cfg_path.clone(),
                     PyvenvCfgParseErrorKind::InvalidHomeValue(io_err),
                 )
             })?;
 
+        let parent_environment = if created_with_uv {
+            parent_environment
+                .and_then(|sys_prefix| {
+                    SysPrefixPath::new(
+                        sys_prefix,
+                        SysPrefixPathOrigin::DerivedFromPyvenvCfg,
+                        system,
+                    )
+                    .and_then(|sys_prefix|Self::new(sys_prefix, system))
+                    .inspect_err(|err| {
+                        tracing::warn!(
+                            "Failed to resolve the parent virtual environment of this ephemeral uv virtual environment \
+                            from the `extends-environment` value specified in the `pyvenv.cfg` file at {pyvenv_cfg_path}. \
+                            Packages installed into the parent environment will not be resolved correctly. \
+                            Underlying error: {err}",
+                        );
+                    })
+                    .ok()
+                })
+                .map(Box::new)
+        } else {
+            None
+        };
+
         // but the `version`/`version_info` key is not read by the standard library,
         // and is provided under different keys depending on which virtual-environment creation tool
         // created the `pyvenv.cfg` file. Lenient parsing is appropriate here:
@@ -218,6 +254,7 @@ impl VirtualEnvironment {
             include_system_site_packages,
             version,
             implementation,
+            parent_environment,
         };
 
         tracing::trace!("Resolved metadata for virtual environment: {metadata:?}");
@@ -230,21 +267,38 @@ impl VirtualEnvironment {
     pub(crate) fn site_packages_directories(
         &self,
         system: &dyn System,
-    ) -> SitePackagesDiscoveryResult<Vec<SystemPathBuf>> {
+    ) -> SitePackagesDiscoveryResult<IndexSet<SystemPathBuf>> {
         let VirtualEnvironment {
             root_path,
             base_executable_home_path,
             include_system_site_packages,
             implementation,
             version,
+            parent_environment,
         } = self;
 
-        let mut site_packages_directories = vec![site_packages_directory_from_sys_prefix(
-            root_path,
-            *version,
-            *implementation,
-            system,
-        )?];
+        let mut site_packages_directories =
+            IndexSet::from([site_packages_directory_from_sys_prefix(
+                root_path,
+                *version,
+                *implementation,
+                system,
+            )?]);
+
+        if let Some(parent_env_site_packages) = parent_environment.as_deref() {
+            match parent_env_site_packages.site_packages_directories(system) {
+                Ok(parent_environment_site_packages) => {
+                    site_packages_directories.extend(parent_environment_site_packages);
+                }
+                Err(err) => {
+                    tracing::warn!(
+                        "Failed to resolve the site-packages directories of this ephemeral uv virtual environment's \
+                        parent environment. Packages installed into the parent environment will not be resolved correctly. \
+                        Underlying error: {err}"
+                    );
+                }
+            }
+        }
 
         if *include_system_site_packages {
             let system_sys_prefix =
@@ -261,7 +315,7 @@ impl VirtualEnvironment {
                     system,
                 ) {
                     Ok(site_packages_directory) => {
-                        site_packages_directories.push(site_packages_directory);
+                        site_packages_directories.insert(site_packages_directory);
                     }
                     Err(error) => tracing::warn!(
                         "{error}. System site-packages will not be used for module resolution."
@@ -309,15 +363,15 @@ impl SystemEnvironment {
     pub(crate) fn site_packages_directories(
         &self,
         system: &dyn System,
-    ) -> SitePackagesDiscoveryResult<Vec<SystemPathBuf>> {
+    ) -> SitePackagesDiscoveryResult<IndexSet<SystemPathBuf>> {
         let SystemEnvironment { root_path } = self;
 
-        let site_packages_directories = vec![site_packages_directory_from_sys_prefix(
+        let site_packages_directories = IndexSet::from([site_packages_directory_from_sys_prefix(
             root_path,
             None,
             PythonImplementation::Unknown,
             system,
-        )?];
+        )?]);
 
         tracing::debug!(
             "Resolved site-packages directories for this environment are: {site_packages_directories:?}"
@@ -550,12 +604,12 @@ impl SysPrefixPath {
         if cfg!(target_os = "windows") {
             Some(Self {
                 inner: path.to_path_buf(),
-                origin: SysPrefixPathOrigin::Derived,
+                origin: SysPrefixPathOrigin::DerivedFromPyvenvCfg,
             })
         } else {
             path.parent().map(|path| Self {
                 inner: path.to_path_buf(),
-                origin: SysPrefixPathOrigin::Derived,
+                origin: SysPrefixPathOrigin::DerivedFromPyvenvCfg,
             })
         }
     }
@@ -575,13 +629,22 @@ impl fmt::Display for SysPrefixPath {
     }
 }
 
+/// Enumeration of sources a `sys.prefix` path can come from.
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub enum SysPrefixPathOrigin {
+    /// The `sys.prefix` path came from a `--python` CLI flag
     PythonCliFlag,
+    /// The `sys.prefix` path came from the `VIRTUAL_ENV` environment variable
     VirtualEnvVar,
+    /// The `sys.prefix` path came from the `CONDA_PREFIX` environment variable
     CondaPrefixVar,
-    Derived,
+    /// The `sys.prefix` path was derived from a value in a `pyvenv.cfg` file:
+    /// either the value associated with the `home` key
+    /// or the value associated with the `extends-environment` key
+    DerivedFromPyvenvCfg,
+    /// A `.venv` directory was found in the current working directory,
+    /// and the `sys.prefix` path is the path to that virtual environment.
     LocalVenv,
 }
 
@@ -591,7 +654,7 @@ impl SysPrefixPathOrigin {
     pub(crate) fn must_be_virtual_env(self) -> bool {
         match self {
             Self::LocalVenv | Self::VirtualEnvVar => true,
-            Self::PythonCliFlag | Self::Derived | Self::CondaPrefixVar => false,
+            Self::PythonCliFlag | Self::DerivedFromPyvenvCfg | Self::CondaPrefixVar => false,
         }
     }
 }
@@ -602,7 +665,7 @@ impl Display for SysPrefixPathOrigin {
             Self::PythonCliFlag => f.write_str("`--python` argument"),
             Self::VirtualEnvVar => f.write_str("`VIRTUAL_ENV` environment variable"),
             Self::CondaPrefixVar => f.write_str("`CONDA_PREFIX` environment variable"),
-            Self::Derived => f.write_str("derived `sys.prefix` path"),
+            Self::DerivedFromPyvenvCfg => f.write_str("derived `sys.prefix` path"),
             Self::LocalVenv => f.write_str("local virtual environment"),
         }
     }
@@ -901,11 +964,14 @@ mod tests {
 
             if self_venv.system_site_packages {
                 assert_eq!(
-                    &site_packages_directories,
+                    site_packages_directories.as_slice(),
                     &[expected_venv_site_packages, expected_system_site_packages]
                 );
             } else {
-                assert_eq!(&site_packages_directories, &[expected_venv_site_packages]);
+                assert_eq!(
+                    site_packages_directories.as_slice(),
+                    &[expected_venv_site_packages]
+                );
             }
         }
 
@@ -946,7 +1012,10 @@ mod tests {
                 ))
             };
 
-            assert_eq!(&site_packages_directories, &[expected_site_packages]);
+            assert_eq!(
+                site_packages_directories.as_slice(),
+                &[expected_site_packages]
+            );
         }
     }