diff --git a/clippy.toml b/clippy.toml index 4ea6eccde0..539d63305b 100644 --- a/clippy.toml +++ b/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 = [ diff --git a/crates/ty_python_semantic/src/site_packages.rs b/crates/ty_python_semantic/src/site_packages.rs index b46dd6cb99..e893c8817e 100644 --- a/crates/ty_python_semantic/src/site_packages.rs +++ b/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) -> Option { + 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, + 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, + implementation: PythonImplementation, system: &dyn System, ) -> SitePackagesDiscoveryResult { 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,15 +800,15 @@ 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!( @@ -743,12 +816,13 @@ mod tests { ); } 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();