Allow discovery of venv in VIRTUAL_ENV env variable (#16853)

MatthewMckee4 · web-flow · commit cdafd8e32b1c · 2025-03-20T13:55:35.000Z
## 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
diff --git 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
@@ -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
@@ -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
@@ -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,11 +143,21 @@ 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.
     ///
     /// This variant is mainly intended for testing where we want to skip resolving `site-packages`
     /// because it would unnecessarily complicate the test setup.
     KnownSitePackages(Vec<SystemPathBuf>),
 }
+
+impl PythonPath {
+    pub fn from_virtual_env_var(path: impl Into<SystemPathBuf>) -> 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
@@ -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<SystemPath>,
+        origin: SysPrefixPathOrigin,
         system: &dyn System,
     ) -> SitePackagesDiscoveryResult<Self> {
-        Self::new_impl(path.as_ref(), system)
+        Self::new_impl(path.as_ref(), origin, system)
     }
 
-    fn new_impl(path: &SystemPath, system: &dyn System) -> SitePackagesDiscoveryResult<Self> {
+    fn new_impl(
+        path: &SystemPath,
+        origin: SysPrefixPathOrigin,
+        system: &dyn System,
+    ) -> SitePackagesDiscoveryResult<Self> {
         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<SystemPath>,
+        origin: SysPrefixPathOrigin,
         system: &dyn System,
     ) -> SitePackagesDiscoveryResult<Self> {
-        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<Self> {
         // 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(
