From cdafd8e32b1c4d76d61d6d6578a8a7fa5a63dab2 Mon Sep 17 00:00:00 2001 From: Matthew Mckee Date: Thu, 20 Mar 2025 13:55:35 +0000 Subject: [PATCH] Allow discovery of venv in VIRTUAL_ENV env variable (#16853) ## Summary Fixes #16744 Allows the cli to find a virtual environment from the VIRTUAL_ENV environment variable if no `--python` is set ## Test Plan Manual testing, of: - Virtual environments explicitly activated using `source .venv/bin/activate` - Virtual environments implicilty activated via `uv run` - Broken virtual environments with no `pyvenv.cfg` file --- crates/red_knot/src/args.rs | 2 + .../red_knot_project/src/metadata/options.rs | 9 +- .../src/module_resolver/resolver.rs | 5 +- .../red_knot_python_semantic/src/program.rs | 13 +- .../src/site_packages.rs | 116 +++++++++++++----- 5 files changed, 111 insertions(+), 34 deletions(-) diff --git a/crates/red_knot/src/args.rs b/crates/red_knot/src/args.rs index 8011bc5cd4..f7863ab53e 100644 --- a/crates/red_knot/src/args.rs +++ b/crates/red_knot/src/args.rs @@ -50,6 +50,8 @@ pub(crate) struct CheckCommand { /// Path to the Python installation from which Red Knot resolves type information and third-party dependencies. /// + /// If not specified, Red Knot will look at the `VIRTUAL_ENV` environment variable. + /// /// Red Knot will search in the path's `site-packages` directories for type information and /// third-party imports. /// diff --git a/crates/red_knot_project/src/metadata/options.rs b/crates/red_knot_project/src/metadata/options.rs index f5a106df43..43a6c4233d 100644 --- a/crates/red_knot_project/src/metadata/options.rs +++ b/crates/red_knot_project/src/metadata/options.rs @@ -106,9 +106,14 @@ impl Options { custom_typeshed: typeshed.map(|path| path.absolute(project_root, system)), python_path: python .map(|python_path| { - PythonPath::SysPrefix(python_path.absolute(project_root, system)) + PythonPath::from_cli_flag(python_path.absolute(project_root, system)) }) - .unwrap_or(PythonPath::KnownSitePackages(vec![])), + .or_else(|| { + std::env::var("VIRTUAL_ENV") + .ok() + .map(PythonPath::from_virtual_env_var) + }) + .unwrap_or_else(|| PythonPath::KnownSitePackages(vec![])), } } diff --git a/crates/red_knot_python_semantic/src/module_resolver/resolver.rs b/crates/red_knot_python_semantic/src/module_resolver/resolver.rs index d983e71895..ba9150f3c9 100644 --- a/crates/red_knot_python_semantic/src/module_resolver/resolver.rs +++ b/crates/red_knot_python_semantic/src/module_resolver/resolver.rs @@ -223,14 +223,15 @@ impl SearchPaths { static_paths.push(stdlib_path); let site_packages_paths = match python_path { - PythonPath::SysPrefix(sys_prefix) => { + PythonPath::SysPrefix(sys_prefix, origin) => { // TODO: We may want to warn here if the venv's python version is older // than the one resolved in the program settings because it indicates // that the `target-version` is incorrectly configured or that the // venv is out of date. - VirtualEnvironment::new(sys_prefix, system) + VirtualEnvironment::new(sys_prefix, *origin, system) .and_then(|venv| venv.site_packages_directories(system))? } + PythonPath::KnownSitePackages(paths) => paths .iter() .map(|path| canonicalize(path, system)) diff --git a/crates/red_knot_python_semantic/src/program.rs b/crates/red_knot_python_semantic/src/program.rs index 4d3c1cc6c6..325d594faf 100644 --- a/crates/red_knot_python_semantic/src/program.rs +++ b/crates/red_knot_python_semantic/src/program.rs @@ -1,5 +1,6 @@ use crate::module_resolver::SearchPaths; use crate::python_platform::PythonPlatform; +use crate::site_packages::SysPrefixPathOrigin; use crate::Db; use anyhow::Context; @@ -142,7 +143,7 @@ pub enum PythonPath { /// `/opt/homebrew/lib/python3.X/site-packages`. /// /// [`sys.prefix`]: https://docs.python.org/3/library/sys.html#sys.prefix - SysPrefix(SystemPathBuf), + SysPrefix(SystemPathBuf, SysPrefixPathOrigin), /// Resolved site packages paths. /// @@ -150,3 +151,13 @@ pub enum PythonPath { /// because it would unnecessarily complicate the test setup. KnownSitePackages(Vec), } + +impl PythonPath { + pub fn from_virtual_env_var(path: impl Into) -> Self { + Self::SysPrefix(path.into(), SysPrefixPathOrigin::VirtualEnvVar) + } + + pub fn from_cli_flag(path: SystemPathBuf) -> Self { + Self::SysPrefix(path, SysPrefixPathOrigin::PythonCliFlag) + } +} diff --git a/crates/red_knot_python_semantic/src/site_packages.rs b/crates/red_knot_python_semantic/src/site_packages.rs index f2812aa70d..6cfbd3ef8a 100644 --- a/crates/red_knot_python_semantic/src/site_packages.rs +++ b/crates/red_knot_python_semantic/src/site_packages.rs @@ -9,6 +9,7 @@ //! on Linux.) use std::fmt; +use std::fmt::Display; use std::io; use std::num::NonZeroUsize; use std::ops::Deref; @@ -43,23 +44,28 @@ pub(crate) struct VirtualEnvironment { impl VirtualEnvironment { pub(crate) fn new( path: impl AsRef, + origin: SysPrefixPathOrigin, system: &dyn System, ) -> SitePackagesDiscoveryResult { - Self::new_impl(path.as_ref(), system) + Self::new_impl(path.as_ref(), origin, system) } - fn new_impl(path: &SystemPath, system: &dyn System) -> SitePackagesDiscoveryResult { + fn new_impl( + path: &SystemPath, + origin: SysPrefixPathOrigin, + system: &dyn System, + ) -> SitePackagesDiscoveryResult { fn pyvenv_cfg_line_number(index: usize) -> NonZeroUsize { index.checked_add(1).and_then(NonZeroUsize::new).unwrap() } - let venv_path = SysPrefixPath::new(path, system)?; + let venv_path = SysPrefixPath::new(path, origin, system)?; let pyvenv_cfg_path = venv_path.join("pyvenv.cfg"); tracing::debug!("Attempting to parse virtual environment metadata at '{pyvenv_cfg_path}'"); let pyvenv_cfg = system .read_to_string(&pyvenv_cfg_path) - .map_err(SitePackagesDiscoveryError::NoPyvenvCfgFile)?; + .map_err(|io_err| SitePackagesDiscoveryError::NoPyvenvCfgFile(origin, io_err))?; let mut include_system_site_packages = false; let mut base_executable_home_path = None; @@ -205,12 +211,12 @@ System site-packages will not be used for module resolution.", #[derive(Debug, thiserror::Error)] pub(crate) enum SitePackagesDiscoveryError { - #[error("Invalid --python argument: `{0}` could not be canonicalized")] - VenvDirCanonicalizationError(SystemPathBuf, #[source] io::Error), - #[error("Invalid --python argument: `{0}` does not point to a directory on disk")] - VenvDirIsNotADirectory(SystemPathBuf), - #[error("--python points to a broken venv with no pyvenv.cfg file")] - NoPyvenvCfgFile(#[source] io::Error), + #[error("Invalid {1}: `{0}` could not be canonicalized")] + VenvDirCanonicalizationError(SystemPathBuf, SysPrefixPathOrigin, #[source] io::Error), + #[error("Invalid {1}: `{0}` does not point to a directory on disk")] + VenvDirIsNotADirectory(SystemPathBuf, SysPrefixPathOrigin), + #[error("{0} points to a broken venv with no pyvenv.cfg file")] + NoPyvenvCfgFile(SysPrefixPathOrigin, #[source] io::Error), #[error("Failed to parse the pyvenv.cfg file at {0} because {1}")] PyvenvCfgParseError(SystemPathBuf, PyvenvCfgParseErrorKind), #[error("Failed to search the `lib` directory of the Python installation at {1} for `site-packages`")] @@ -370,18 +376,23 @@ fn site_packages_directory_from_sys_prefix( /// /// [`sys.prefix`]: https://docs.python.org/3/library/sys.html#sys.prefix #[derive(Debug, PartialEq, Eq, Clone)] -pub(crate) struct SysPrefixPath(SystemPathBuf); +pub(crate) struct SysPrefixPath { + inner: SystemPathBuf, + origin: SysPrefixPathOrigin, +} impl SysPrefixPath { fn new( unvalidated_path: impl AsRef, + origin: SysPrefixPathOrigin, system: &dyn System, ) -> SitePackagesDiscoveryResult { - Self::new_impl(unvalidated_path.as_ref(), system) + Self::new_impl(unvalidated_path.as_ref(), origin, system) } fn new_impl( unvalidated_path: &SystemPath, + origin: SysPrefixPathOrigin, system: &dyn System, ) -> SitePackagesDiscoveryResult { // It's important to resolve symlinks here rather than simply making the path absolute, @@ -392,14 +403,21 @@ impl SysPrefixPath { .map_err(|io_err| { SitePackagesDiscoveryError::VenvDirCanonicalizationError( unvalidated_path.to_path_buf(), + origin, io_err, ) })?; system .is_directory(&canonicalized) - .then_some(Self(canonicalized)) + .then_some(Self { + inner: canonicalized, + origin, + }) .ok_or_else(|| { - SitePackagesDiscoveryError::VenvDirIsNotADirectory(unvalidated_path.to_path_buf()) + SitePackagesDiscoveryError::VenvDirIsNotADirectory( + unvalidated_path.to_path_buf(), + origin, + ) }) } @@ -408,9 +426,15 @@ impl SysPrefixPath { // the parent of a canonicalised path that is known to exist // is guaranteed to be a directory. if cfg!(target_os = "windows") { - Some(Self(path.to_path_buf())) + Some(Self { + inner: path.to_path_buf(), + origin: SysPrefixPathOrigin::Derived, + }) } else { - path.parent().map(|path| Self(path.to_path_buf())) + path.parent().map(|path| Self { + inner: path.to_path_buf(), + origin: SysPrefixPathOrigin::Derived, + }) } } } @@ -419,13 +443,31 @@ impl Deref for SysPrefixPath { type Target = SystemPath; fn deref(&self) -> &Self::Target { - &self.0 + &self.inner } } impl fmt::Display for SysPrefixPath { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "`sys.prefix` path `{}`", self.0) + write!(f, "`sys.prefix` path `{}`", self.inner) + } +} + +#[derive(Debug, PartialEq, Eq, Copy, Clone)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +pub enum SysPrefixPathOrigin { + PythonCliFlag, + VirtualEnvVar, + Derived, +} + +impl Display for SysPrefixPathOrigin { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match self { + Self::PythonCliFlag => f.write_str("`--python` argument"), + Self::VirtualEnvVar => f.write_str("`VIRTUAL_ENV` environment variable"), + Self::Derived => f.write_str("derived `sys.prefix` path"), + } } } @@ -584,11 +626,19 @@ mod tests { fn test(self) { let venv_path = self.build_mock_venv(); - let venv = VirtualEnvironment::new(venv_path.clone(), &self.system).unwrap(); + let venv = VirtualEnvironment::new( + venv_path.clone(), + SysPrefixPathOrigin::VirtualEnvVar, + &self.system, + ) + .unwrap(); assert_eq!( venv.venv_path, - SysPrefixPath(self.system.canonicalize_path(&venv_path).unwrap()) + SysPrefixPath { + inner: self.system.canonicalize_path(&venv_path).unwrap(), + origin: SysPrefixPathOrigin::VirtualEnvVar, + } ); assert_eq!(venv.include_system_site_packages, self.system_site_packages); @@ -730,7 +780,7 @@ mod tests { fn reject_venv_that_does_not_exist() { let system = TestSystem::default(); assert!(matches!( - VirtualEnvironment::new("/.venv", &system), + VirtualEnvironment::new("/.venv", SysPrefixPathOrigin::VirtualEnvVar, &system), Err(SitePackagesDiscoveryError::VenvDirCanonicalizationError(..)) )); } @@ -743,7 +793,7 @@ mod tests { .write_file_all("/.venv", "") .unwrap(); assert!(matches!( - VirtualEnvironment::new("/.venv", &system), + VirtualEnvironment::new("/.venv", SysPrefixPathOrigin::VirtualEnvVar, &system), Err(SitePackagesDiscoveryError::VenvDirIsNotADirectory(..)) )); } @@ -756,8 +806,11 @@ mod tests { .create_directory_all("/.venv") .unwrap(); assert!(matches!( - VirtualEnvironment::new("/.venv", &system), - Err(SitePackagesDiscoveryError::NoPyvenvCfgFile(_)) + VirtualEnvironment::new("/.venv", SysPrefixPathOrigin::VirtualEnvVar, &system), + Err(SitePackagesDiscoveryError::NoPyvenvCfgFile( + SysPrefixPathOrigin::VirtualEnvVar, + _ + )) )); } @@ -769,7 +822,8 @@ mod tests { memory_fs .write_file_all(&pyvenv_cfg_path, "home = bar = /.venv/bin") .unwrap(); - let venv_result = VirtualEnvironment::new("/.venv", &system); + let venv_result = + VirtualEnvironment::new("/.venv", SysPrefixPathOrigin::VirtualEnvVar, &system); assert!(matches!( venv_result, Err(SitePackagesDiscoveryError::PyvenvCfgParseError( @@ -788,7 +842,8 @@ mod tests { memory_fs .write_file_all(&pyvenv_cfg_path, "home =") .unwrap(); - let venv_result = VirtualEnvironment::new("/.venv", &system); + let venv_result = + VirtualEnvironment::new("/.venv", SysPrefixPathOrigin::VirtualEnvVar, &system); assert!(matches!( venv_result, Err(SitePackagesDiscoveryError::PyvenvCfgParseError( @@ -807,7 +862,8 @@ mod tests { memory_fs .write_file_all(&pyvenv_cfg_path, "= whatever") .unwrap(); - let venv_result = VirtualEnvironment::new("/.venv", &system); + let venv_result = + VirtualEnvironment::new("/.venv", SysPrefixPathOrigin::VirtualEnvVar, &system); assert!(matches!( venv_result, Err(SitePackagesDiscoveryError::PyvenvCfgParseError( @@ -824,7 +880,8 @@ mod tests { let memory_fs = system.memory_file_system(); let pyvenv_cfg_path = SystemPathBuf::from("/.venv/pyvenv.cfg"); memory_fs.write_file_all(&pyvenv_cfg_path, "").unwrap(); - let venv_result = VirtualEnvironment::new("/.venv", &system); + let venv_result = + VirtualEnvironment::new("/.venv", SysPrefixPathOrigin::VirtualEnvVar, &system); assert!(matches!( venv_result, Err(SitePackagesDiscoveryError::PyvenvCfgParseError( @@ -843,7 +900,8 @@ mod tests { memory_fs .write_file_all(&pyvenv_cfg_path, "home = foo") .unwrap(); - let venv_result = VirtualEnvironment::new("/.venv", &system); + let venv_result = + VirtualEnvironment::new("/.venv", SysPrefixPathOrigin::VirtualEnvVar, &system); assert!(matches!( venv_result, Err(SitePackagesDiscoveryError::PyvenvCfgParseError(