[ty] Add support for PyPy virtual environments (astral-sh#18203)

Co-authored-by: Alex Waygood <[email protected]>

Author: Mathemmagician

clippy.toml

@@ -1,6 +1,7 @@
 doc-valid-idents = [
     "..",
     "CodeQL",
+    "CPython",
     "FastAPI",
     "IPython",
     "LangChain",
@@ -14,7 +15,7 @@ doc-valid-idents = [
     "SNMPv1",
     "SNMPv2",
     "SNMPv3",
-    "PyFlakes"
+    "PyFlakes",
 ]
 
 ignore-interior-mutability = [

crates/ty_python_semantic/src/site_packages.rs

@@ -62,6 +62,34 @@ impl PythonEnvironment {
     }
 }
 
+/// The Python runtime that produced the venv.
+///
+/// We only need to distinguish cases that change the on-disk layout.
+/// Everything else can be treated like CPython.
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub(crate) enum PythonImplementation {
+    CPython,
+    PyPy,
+    GraalPy,
+    /// Fallback when the value is missing or unrecognised.
+    /// We treat it like CPython but keep the information for diagnostics.
+    Unknown,
+}
+
+impl PythonImplementation {
+    /// Return the relative path from `sys.prefix` to the `site-packages` directory
+    /// if this is a known implementation. Return `None` if this is an unknown implementation.
+    fn relative_site_packages_path(self, version: Option<PythonVersion>) -> Option<String> {
+        match self {
+            Self::CPython | Self::GraalPy => {
+                version.map(|version| format!("lib/python{version}/site-packages"))
+            }
+            Self::PyPy => version.map(|version| format!("lib/pypy{version}/site-packages")),
+            Self::Unknown => None,
+        }
+    }
+}
+
 /// Abstraction for a Python virtual environment.
 ///
 /// Most of this information is derived from the virtual environment's `pyvenv.cfg` file.
@@ -82,6 +110,7 @@ pub(crate) struct VirtualEnvironment {
     /// in an acceptable format under any of the keys we expect.
     /// This field will be `None` if so.
     version: Option<PythonVersion>,
+    implementation: PythonImplementation,
 }
 
 impl VirtualEnvironment {
@@ -104,6 +133,7 @@ impl VirtualEnvironment {
         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;
 
         // 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
@@ -140,6 +170,14 @@ impl VirtualEnvironment {
                     // `virtualenv` and `uv` call this key `version_info`,
                     // but the stdlib venv module calls it `version`
                     "version" | "version_info" => version_info_string = Some(value),
+                    "implementation" => {
+                        implementation = match value.to_ascii_lowercase().as_str() {
+                            "cpython" => PythonImplementation::CPython,
+                            "graalvm" => PythonImplementation::GraalPy,
+                            "pypy" => PythonImplementation::PyPy,
+                            _ => PythonImplementation::Unknown,
+                        };
+                    }
                     _ => continue,
                 }
             }
@@ -179,6 +217,7 @@ impl VirtualEnvironment {
             base_executable_home_path,
             include_system_site_packages,
             version,
+            implementation,
         };
 
         tracing::trace!("Resolved metadata for virtual environment: {metadata:?}");
@@ -196,11 +235,15 @@ impl VirtualEnvironment {
             root_path,
             base_executable_home_path,
             include_system_site_packages,
+            implementation,
             version,
         } = self;
 
         let mut site_packages_directories = vec![site_packages_directory_from_sys_prefix(
-            root_path, *version, system,
+            root_path,
+            *version,
+            *implementation,
+            system,
         )?];
 
         if *include_system_site_packages {
@@ -211,7 +254,12 @@ impl VirtualEnvironment {
             // or if we fail to resolve the `site-packages` from the `sys.prefix` path,
             // we should probably print a warning but *not* abort type checking
             if let Some(sys_prefix_path) = system_sys_prefix {
-                match site_packages_directory_from_sys_prefix(&sys_prefix_path, *version, system) {
+                match site_packages_directory_from_sys_prefix(
+                    &sys_prefix_path,
+                    *version,
+                    *implementation,
+                    system,
+                ) {
                     Ok(site_packages_directory) => {
                         site_packages_directories.push(site_packages_directory);
                     }
@@ -265,7 +313,10 @@ impl SystemEnvironment {
         let SystemEnvironment { root_path } = self;
 
         let site_packages_directories = vec![site_packages_directory_from_sys_prefix(
-            root_path, None, system,
+            root_path,
+            None,
+            PythonImplementation::Unknown,
+            system,
         )?];
 
         tracing::debug!(
@@ -330,6 +381,7 @@ when trying to resolve the `home` value to a directory on disk: {io_err}"
 fn site_packages_directory_from_sys_prefix(
     sys_prefix_path: &SysPrefixPath,
     python_version: Option<PythonVersion>,
+    implementation: PythonImplementation,
     system: &dyn System,
 ) -> SitePackagesDiscoveryResult<SystemPathBuf> {
     tracing::debug!("Searching for site-packages directory in {sys_prefix_path}");
@@ -369,15 +421,21 @@ fn site_packages_directory_from_sys_prefix(
 
     // If we were able to figure out what Python version this installation is,
     // we should be able to avoid iterating through all items in the `lib/` directory:
-    if let Some(version) = python_version {
-        let expected_path = sys_prefix_path.join(format!("lib/python{version}/site-packages"));
-        if system.is_directory(&expected_path) {
-            return Ok(expected_path);
+    if let Some(expected_relative_path) = implementation.relative_site_packages_path(python_version)
+    {
+        let expected_absolute_path = sys_prefix_path.join(expected_relative_path);
+        if system.is_directory(&expected_absolute_path) {
+            return Ok(expected_absolute_path);
         }
-        if version.free_threaded_build_available() {
-            // Nearly the same as `expected_path`, but with an additional `t` after {version}:
-            let alternative_path =
-                sys_prefix_path.join(format!("lib/python{version}t/site-packages"));
+
+        // CPython free-threaded (3.13+) variant: pythonXYt
+        if matches!(implementation, PythonImplementation::CPython)
+            && python_version.is_some_and(PythonVersion::free_threaded_build_available)
+        {
+            let alternative_path = sys_prefix_path.join(format!(
+                "lib/python{}t/site-packages",
+                python_version.unwrap()
+            ));
             if system.is_directory(&alternative_path) {
                 return Ok(alternative_path);
             }
@@ -412,7 +470,7 @@ fn site_packages_directory_from_sys_prefix(
             .file_name()
             .expect("File name to be non-null because path is guaranteed to be a child of `lib`");
 
-        if !name.starts_with("python3.") {
+        if !(name.starts_with("python3.") || name.starts_with("pypy3.")) {
             continue;
         }
 
@@ -623,10 +681,20 @@ mod tests {
 
     use super::*;
 
+    impl PythonEnvironment {
+        fn expect_venv(self) -> VirtualEnvironment {
+            match self {
+                Self::Virtual(venv) => venv,
+                Self::System(_) => panic!("Expected a virtual environment"),
+            }
+        }
+    }
+
     struct VirtualEnvironmentTestCase {
         system_site_packages: bool,
         pyvenv_cfg_version_field: Option<&'static str>,
         command_field: Option<&'static str>,
+        implementation_field: Option<&'static str>,
     }
 
     struct PythonEnvironmentTestCase {
@@ -679,6 +747,7 @@ mod tests {
                 pyvenv_cfg_version_field,
                 system_site_packages,
                 command_field,
+                implementation_field,
             }) = virtual_env
             else {
                 return system_install_sys_prefix;
@@ -709,6 +778,10 @@ mod tests {
                 pyvenv_cfg_contents.push_str(command_field);
                 pyvenv_cfg_contents.push('\n');
             }
+            if let Some(implementation_field) = implementation_field {
+                pyvenv_cfg_contents.push_str(implementation_field);
+                pyvenv_cfg_contents.push('\n');
+            }
             // Deliberately using weird casing here to test that our pyvenv.cfg parsing is case-insensitive:
             if *system_site_packages {
                 pyvenv_cfg_contents.push_str("include-system-site-packages = TRuE\n");
@@ -727,28 +800,29 @@ mod tests {
         }
 
         #[track_caller]
-        fn run(self) {
+        fn run(self) -> PythonEnvironment {
             let env_path = self.build();
             let env = PythonEnvironment::new(env_path.clone(), self.origin, &self.system)
                 .expect("Expected environment construction to succeed");
 
             let expect_virtual_env = self.virtual_env.is_some();
-            match env {
+            match &env {
                 PythonEnvironment::Virtual(venv) if expect_virtual_env => {
-                    self.assert_virtual_environment(&venv, &env_path);
+                    self.assert_virtual_environment(venv, &env_path);
                 }
                 PythonEnvironment::Virtual(venv) => {
                     panic!(
                         "Expected a system environment, but got a virtual environment: {venv:?}"
                     );
                 }
                 PythonEnvironment::System(env) if !expect_virtual_env => {
-                    self.assert_system_environment(&env, &env_path);
+                    self.assert_system_environment(env, &env_path);
                 }
                 PythonEnvironment::System(env) => {
                     panic!("Expected a virtual environment, but got a system environment: {env:?}");
                 }
             }
+            env
         }
 
         fn assert_virtual_environment(
@@ -941,6 +1015,7 @@ mod tests {
                 system_site_packages: false,
                 pyvenv_cfg_version_field: None,
                 command_field: None,
+                implementation_field: None,
             }),
         };
         test.run();
@@ -957,6 +1032,7 @@ mod tests {
                 system_site_packages: false,
                 pyvenv_cfg_version_field: Some("version = 3.12"),
                 command_field: None,
+                implementation_field: None,
             }),
         };
         test.run();
@@ -973,6 +1049,7 @@ mod tests {
                 system_site_packages: false,
                 pyvenv_cfg_version_field: Some("version_info = 3.12"),
                 command_field: None,
+                implementation_field: None,
             }),
         };
         test.run();
@@ -989,6 +1066,7 @@ mod tests {
                 system_site_packages: false,
                 pyvenv_cfg_version_field: Some("version_info = 3.12.0rc2"),
                 command_field: None,
+                implementation_field: None,
             }),
         };
         test.run();
@@ -1005,6 +1083,7 @@ mod tests {
                 system_site_packages: false,
                 pyvenv_cfg_version_field: Some("version_info = 3.13"),
                 command_field: None,
+                implementation_field: None,
             }),
         };
         test.run();
@@ -1021,11 +1100,84 @@ mod tests {
                 system_site_packages: true,
                 pyvenv_cfg_version_field: Some("version_info = 3.13"),
                 command_field: None,
+                implementation_field: None,
             }),
         };
         test.run();
     }
 
+    #[test]
+    fn detects_pypy_implementation() {
+        let test = PythonEnvironmentTestCase {
+            system: TestSystem::default(),
+            minor_version: 13,
+            free_threaded: true,
+            origin: SysPrefixPathOrigin::VirtualEnvVar,
+            virtual_env: Some(VirtualEnvironmentTestCase {
+                system_site_packages: true,
+                pyvenv_cfg_version_field: None,
+                command_field: None,
+                implementation_field: Some("implementation = PyPy"),
+            }),
+        };
+        let venv = test.run().expect_venv();
+        assert_eq!(venv.implementation, PythonImplementation::PyPy);
+    }
+
+    #[test]
+    fn detects_cpython_implementation() {
+        let test = PythonEnvironmentTestCase {
+            system: TestSystem::default(),
+            minor_version: 13,
+            free_threaded: true,
+            origin: SysPrefixPathOrigin::VirtualEnvVar,
+            virtual_env: Some(VirtualEnvironmentTestCase {
+                system_site_packages: true,
+                pyvenv_cfg_version_field: None,
+                command_field: None,
+                implementation_field: Some("implementation = CPython"),
+            }),
+        };
+        let venv = test.run().expect_venv();
+        assert_eq!(venv.implementation, PythonImplementation::CPython);
+    }
+
+    #[test]
+    fn detects_graalpy_implementation() {
+        let test = PythonEnvironmentTestCase {
+            system: TestSystem::default(),
+            minor_version: 13,
+            free_threaded: true,
+            origin: SysPrefixPathOrigin::VirtualEnvVar,
+            virtual_env: Some(VirtualEnvironmentTestCase {
+                system_site_packages: true,
+                pyvenv_cfg_version_field: None,
+                command_field: None,
+                implementation_field: Some("implementation = GraalVM"),
+            }),
+        };
+        let venv = test.run().expect_venv();
+        assert_eq!(venv.implementation, PythonImplementation::GraalPy);
+    }
+
+    #[test]
+    fn detects_unknown_implementation() {
+        let test = PythonEnvironmentTestCase {
+            system: TestSystem::default(),
+            minor_version: 13,
+            free_threaded: true,
+            origin: SysPrefixPathOrigin::VirtualEnvVar,
+            virtual_env: Some(VirtualEnvironmentTestCase {
+                system_site_packages: true,
+                pyvenv_cfg_version_field: None,
+                command_field: None,
+                implementation_field: None,
+            }),
+        };
+        let venv = test.run().expect_venv();
+        assert_eq!(venv.implementation, PythonImplementation::Unknown);
+    }
+
     #[test]
     fn reject_env_that_does_not_exist() {
         let system = TestSystem::default();
@@ -1122,6 +1274,7 @@ mod tests {
                 command_field: Some(
                     r#"command = /.pyenv/versions/3.13.3/bin/python3.13 -m venv --without-pip --prompt="python-default/3.13.3" /somewhere-else/python/virtualenvs/python-default/3.13.3"#,
                 ),
+                implementation_field: None,
             }),
         };
         test.run();