[ty] Resolve python environment in Options::to_program_settings (#18960)

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
This commit is contained in:
Micha Reiser 2025-06-26 17:57:16 +02:00 committed by GitHub
parent d00697621e
commit 1dcdf7f41d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 208 additions and 277 deletions

View file

@ -11,12 +11,12 @@ pub use module_resolver::{
system_module_search_paths,
};
pub use program::{
Program, ProgramSettings, PythonEnvironmentPath, PythonVersionFileSource, PythonVersionSource,
Program, ProgramSettings, PythonVersionFileSource, PythonVersionSource,
PythonVersionWithSource, SearchPathSettings,
};
pub use python_platform::PythonPlatform;
pub use semantic_model::{HasType, SemanticModel};
pub use site_packages::SysPrefixPathOrigin;
pub use site_packages::{PythonEnvironment, SitePackagesPaths, SysPrefixPathOrigin};
pub use util::diagnostics::add_inferred_python_version_hint_to_diagnostic;
pub mod ast_node_ref;

View file

@ -1,9 +1,8 @@
use std::borrow::Cow;
use std::fmt;
use std::iter::FusedIterator;
use std::str::{FromStr, Split};
use std::str::Split;
use camino::Utf8Component;
use compact_str::format_compact;
use rustc_hash::{FxBuildHasher, FxHashSet};
@ -15,13 +14,7 @@ use ruff_python_ast::PythonVersion;
use crate::db::Db;
use crate::module_name::ModuleName;
use crate::module_resolver::typeshed::{TypeshedVersions, vendored_typeshed_versions};
use crate::site_packages::{
PythonEnvironment, SitePackagesDiscoveryError, SitePackagesPaths, SysPrefixPathOrigin,
};
use crate::{
Program, PythonEnvironmentPath, PythonVersionSource, PythonVersionWithSource,
SearchPathSettings,
};
use crate::{Program, SearchPathSettings};
use super::module::{Module, ModuleKind};
use super::path::{ModulePath, SearchPath, SearchPathValidationError};
@ -160,13 +153,6 @@ pub struct SearchPaths {
site_packages: Vec<SearchPath>,
typeshed_versions: TypeshedVersions,
/// The Python version implied by the virtual environment.
///
/// If this environment was a system installation or the `pyvenv.cfg` file
/// of the virtual environment did not contain a `version` or `version_info` key,
/// this field will be `None`.
python_version_from_pyvenv_cfg: Option<PythonVersionWithSource>,
}
impl SearchPaths {
@ -191,7 +177,7 @@ impl SearchPaths {
extra_paths,
src_roots,
custom_typeshed: typeshed,
python_environment: python_path,
site_packages_paths,
} = settings;
let mut static_paths = vec![];
@ -236,30 +222,11 @@ impl SearchPaths {
static_paths.push(stdlib_path);
let (site_packages_paths, python_version) = match python_path {
PythonEnvironmentPath::Discover(project_root) => {
Self::discover_python_environment(system, project_root)?
}
PythonEnvironmentPath::Explicit(prefix, origin) => {
tracing::debug!("Resolving {origin}: {prefix}");
PythonEnvironment::new(prefix, origin.clone(), system)?.into_settings(system)?
}
PythonEnvironmentPath::Testing(paths) => (
paths
.iter()
.map(|path| canonicalize(path, system))
.collect(),
None,
),
};
let mut site_packages: Vec<_> = Vec::with_capacity(site_packages_paths.len());
for path in site_packages_paths {
tracing::debug!("Adding site-packages search path '{path}'");
site_packages.push(SearchPath::site_packages(system, path)?);
site_packages.push(SearchPath::site_packages(system, path.clone())?);
}
// TODO vendor typeshed's third-party stubs as well as the stdlib and fallback to them as a final step
@ -285,68 +252,9 @@ impl SearchPaths {
static_paths,
site_packages,
typeshed_versions,
python_version_from_pyvenv_cfg: python_version,
})
}
fn discover_python_environment(
system: &dyn System,
project_root: &SystemPath,
) -> Result<(SitePackagesPaths, Option<PythonVersionWithSource>), SitePackagesDiscoveryError>
{
fn resolve_environment(
system: &dyn System,
path: &SystemPath,
origin: SysPrefixPathOrigin,
) -> Result<(SitePackagesPaths, Option<PythonVersionWithSource>), SitePackagesDiscoveryError>
{
tracing::debug!("Resolving {origin}: {path}");
PythonEnvironment::new(path, origin, system)?.into_settings(system)
}
if let Ok(virtual_env) = system.env_var("VIRTUAL_ENV") {
return resolve_environment(
system,
SystemPath::new(&virtual_env),
SysPrefixPathOrigin::VirtualEnvVar,
);
}
if let Ok(conda_env) = system.env_var("CONDA_PREFIX") {
return resolve_environment(
system,
SystemPath::new(&conda_env),
SysPrefixPathOrigin::CondaPrefixVar,
);
}
tracing::debug!("Discovering virtual environment in `{project_root}`");
let virtual_env_directory = project_root.join(".venv");
match PythonEnvironment::new(
&virtual_env_directory,
SysPrefixPathOrigin::LocalVenv,
system,
)
.and_then(|venv| venv.into_settings(system))
{
Ok(settings) => return Ok(settings),
Err(err) => {
if system.is_directory(&virtual_env_directory) {
tracing::debug!(
"Ignoring automatically detected virtual environment at `{}`: {}",
&virtual_env_directory,
err
);
}
}
}
tracing::debug!("No virtual environment found");
Ok((SitePackagesPaths::default(), None))
}
pub(crate) fn try_register_static_roots(&self, db: &dyn Db) {
let files = db.files();
for path in self.static_paths.iter().chain(self.site_packages.iter()) {
@ -379,52 +287,6 @@ impl SearchPaths {
pub(crate) fn typeshed_versions(&self) -> &TypeshedVersions {
&self.typeshed_versions
}
pub fn try_resolve_installation_python_version(&self) -> Option<Cow<PythonVersionWithSource>> {
if let Some(version) = self.python_version_from_pyvenv_cfg.as_ref() {
return Some(Cow::Borrowed(version));
}
if cfg!(windows) {
// The path to `site-packages` on Unix is
// `<sys.prefix>/lib/pythonX.Y/site-packages`,
// but on Windows it's `<sys.prefix>/Lib/site-packages`.
return None;
}
let primary_site_packages = self.site_packages.first()?.as_system_path()?;
let mut site_packages_ancestor_components =
primary_site_packages.components().rev().skip(1).map(|c| {
// This should have all been validated in `site_packages.rs`
// when we resolved the search paths for the project.
debug_assert!(
matches!(c, Utf8Component::Normal(_)),
"Unexpected component in site-packages path `{c:?}` \
(expected `site-packages` to be an absolute path with symlinks resolved, \
located at `<sys.prefix>/lib/pythonX.Y/site-packages`)"
);
c.as_str()
});
let parent_component = site_packages_ancestor_components.next()?;
if site_packages_ancestor_components.next()? != "lib" {
return None;
}
let version = parent_component
.strip_prefix("python")
.or_else(|| parent_component.strip_prefix("pypy"))?
.trim_end_matches('t');
let version = PythonVersion::from_str(version).ok()?;
let source = PythonVersionSource::InstallationDirectoryLayout {
site_packages_parent_dir: Box::from(parent_component),
};
Some(Cow::Owned(PythonVersionWithSource { version, source }))
}
}
/// Collect all dynamic search paths. For each `site-packages` path:
@ -443,7 +305,6 @@ pub(crate) fn dynamic_resolution_paths(db: &dyn Db) -> Vec<SearchPath> {
static_paths,
site_packages,
typeshed_versions: _,
python_version_from_pyvenv_cfg: _,
} = Program::get(db).search_paths(db);
let mut dynamic_paths = Vec::new();
@ -1534,7 +1395,7 @@ mod tests {
python_platform: PythonPlatform::default(),
search_paths: SearchPathSettings {
custom_typeshed: Some(custom_typeshed),
python_environment: PythonEnvironmentPath::Testing(vec![site_packages]),
site_packages_paths: vec![site_packages],
..SearchPathSettings::new(vec![src.clone()])
}
.to_search_paths(db.system(), db.vendored())
@ -2049,10 +1910,7 @@ not_a_directory
python_version: PythonVersionWithSource::default(),
python_platform: PythonPlatform::default(),
search_paths: SearchPathSettings {
python_environment: PythonEnvironmentPath::Testing(vec![
venv_site_packages,
system_site_packages,
]),
site_packages_paths: vec![venv_site_packages, system_site_packages],
..SearchPathSettings::new(vec![SystemPathBuf::from("/src")])
}
.to_search_paths(db.system(), db.vendored())
@ -2166,7 +2024,7 @@ not_a_directory
python_version: PythonVersionWithSource::default(),
python_platform: PythonPlatform::default(),
search_paths: SearchPathSettings {
python_environment: PythonEnvironmentPath::Testing(vec![site_packages.clone()]),
site_packages_paths: vec![site_packages.clone()],
..SearchPathSettings::new(vec![project_directory])
}
.to_search_paths(db.system(), db.vendored())

View file

@ -7,10 +7,7 @@ use ruff_python_ast::PythonVersion;
use crate::db::tests::TestDb;
use crate::program::{Program, SearchPathSettings};
use crate::{
ProgramSettings, PythonEnvironmentPath, PythonPlatform, PythonVersionSource,
PythonVersionWithSource,
};
use crate::{ProgramSettings, PythonPlatform, PythonVersionSource, PythonVersionWithSource};
/// A test case for the module resolver.
///
@ -246,7 +243,7 @@ impl TestCaseBuilder<MockedTypeshed> {
python_platform,
search_paths: SearchPathSettings {
custom_typeshed: Some(typeshed.clone()),
python_environment: PythonEnvironmentPath::Testing(vec![site_packages.clone()]),
site_packages_paths: vec![site_packages.clone()],
..SearchPathSettings::new(vec![src.clone()])
}
.to_search_paths(db.system(), db.vendored())
@ -306,7 +303,7 @@ impl TestCaseBuilder<VendoredTypeshed> {
},
python_platform,
search_paths: SearchPathSettings {
python_environment: PythonEnvironmentPath::Testing(vec![site_packages.clone()]),
site_packages_paths: vec![site_packages.clone()],
..SearchPathSettings::new(vec![src.clone()])
}
.to_search_paths(db.system(), db.vendored())

View file

@ -1,8 +1,8 @@
use std::sync::Arc;
use crate::Db;
use crate::module_resolver::{SearchPathValidationError, SearchPaths};
use crate::python_platform::PythonPlatform;
use crate::{Db, SysPrefixPathOrigin};
use ruff_db::diagnostic::Span;
use ruff_db::files::system_path_to_file;
@ -173,9 +173,8 @@ pub struct SearchPathSettings {
/// bundled as a zip file in the binary
pub custom_typeshed: Option<SystemPathBuf>,
/// Path to the Python environment from which ty resolves third party dependencies
/// and their type information.
pub python_environment: PythonEnvironmentPath,
/// List of site packages paths to use.
pub site_packages_paths: Vec<SystemPathBuf>,
}
impl SearchPathSettings {
@ -191,7 +190,7 @@ impl SearchPathSettings {
src_roots: vec![],
extra_paths: vec![],
custom_typeshed: None,
python_environment: PythonEnvironmentPath::Testing(vec![]),
site_packages_paths: vec![],
}
}
@ -203,43 +202,3 @@ impl SearchPathSettings {
SearchPaths::from_settings(self, system, vendored)
}
}
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum PythonEnvironmentPath {
/// The path to the Python environment isn't known. Try to discover the Python environment
/// by inspecting environment variables, the project structure, etc. and derive the path from it.
///
/// The path is the project root in which to search for a Python environment.
Discover(SystemPathBuf),
/// Path to a Python environment that is explicitly specified.
///
/// The path that either represents the value of [`sys.prefix`] at runtime in Python
/// for a given Python executable, or which represents a path relative to `sys.prefix`
/// that we will attempt later to resolve into `sys.prefix`. Exactly which this variant
/// represents depends on the [`SysPrefixPathOrigin`] element in the tuple.
///
/// For the case of a virtual environment, where a
/// Python binary is at `/.venv/bin/python`, `sys.prefix` is the path to
/// the virtual environment the Python binary lies inside, i.e. `/.venv`,
/// and `site-packages` will be at `.venv/lib/python3.X/site-packages`.
/// System Python installations generally work the same way: if a system
/// Python installation lies at `/opt/homebrew/bin/python`, `sys.prefix`
/// will be `/opt/homebrew`, and `site-packages` will be at
/// `/opt/homebrew/lib/python3.X/site-packages`.
///
/// [`sys.prefix`]: https://docs.python.org/3/library/sys.html#sys.prefix
Explicit(SystemPathBuf, SysPrefixPathOrigin),
/// Don't search for a Python environment, instead use the provided 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.
Testing(Vec<SystemPathBuf>),
}
impl PythonEnvironmentPath {
pub fn explicit(path: impl Into<SystemPathBuf>, origin: SysPrefixPathOrigin) -> Self {
Self::Explicit(path.into(), origin)
}
}

View file

@ -11,8 +11,11 @@
use std::io;
use std::num::NonZeroUsize;
use std::ops::Deref;
use std::str::FromStr;
use std::{fmt, sync::Arc};
use crate::{PythonVersionFileSource, PythonVersionSource, PythonVersionWithSource};
use camino::Utf8Component;
use indexmap::IndexSet;
use ruff_annotate_snippets::{Level, Renderer, Snippet};
use ruff_db::system::{System, SystemPath, SystemPathBuf};
@ -21,8 +24,6 @@ use ruff_python_trivia::Cursor;
use ruff_source_file::{LineIndex, OneIndexed, SourceCode};
use ruff_text_size::{TextLen, TextRange};
use crate::{PythonVersionFileSource, PythonVersionSource, PythonVersionWithSource};
type SitePackagesDiscoveryResult<T> = Result<T, SitePackagesDiscoveryError>;
/// An ordered, deduplicated set of `site-packages` search paths.
@ -43,13 +44,9 @@ type SitePackagesDiscoveryResult<T> = Result<T, SitePackagesDiscoveryError>;
/// *might* be added to the `SitePackagesPaths` twice, but we wouldn't
/// want duplicates to appear in this set.
#[derive(Debug, PartialEq, Eq, Default)]
pub(crate) struct SitePackagesPaths(IndexSet<SystemPathBuf>);
pub struct SitePackagesPaths(IndexSet<SystemPathBuf>);
impl SitePackagesPaths {
pub(crate) fn len(&self) -> usize {
self.0.len()
}
fn single(path: SystemPathBuf) -> Self {
Self(IndexSet::from([path]))
}
@ -61,6 +58,54 @@ impl SitePackagesPaths {
fn extend(&mut self, other: Self) {
self.0.extend(other.0);
}
/// Tries to detect the version from the layout of the `site-packages` directory.
pub fn python_version_from_layout(&self) -> Option<PythonVersionWithSource> {
if cfg!(windows) {
// The path to `site-packages` on Unix is
// `<sys.prefix>/lib/pythonX.Y/site-packages`,
// but on Windows it's `<sys.prefix>/Lib/site-packages`.
return None;
}
let primary_site_packages = self.0.first()?;
let mut site_packages_ancestor_components =
primary_site_packages.components().rev().skip(1).map(|c| {
// This should have all been validated in `site_packages.rs`
// when we resolved the search paths for the project.
debug_assert!(
matches!(c, Utf8Component::Normal(_)),
"Unexpected component in site-packages path `{c:?}` \
(expected `site-packages` to be an absolute path with symlinks resolved, \
located at `<sys.prefix>/lib/pythonX.Y/site-packages`)"
);
c.as_str()
});
let parent_component = site_packages_ancestor_components.next()?;
if site_packages_ancestor_components.next()? != "lib" {
return None;
}
let version = parent_component
.strip_prefix("python")
.or_else(|| parent_component.strip_prefix("pypy"))?
.trim_end_matches('t');
let version = PythonVersion::from_str(version).ok()?;
let source = PythonVersionSource::InstallationDirectoryLayout {
site_packages_parent_dir: Box::from(parent_component),
};
Some(PythonVersionWithSource { version, source })
}
pub fn into_vec(self) -> Vec<SystemPathBuf> {
self.0.into_iter().collect()
}
}
impl FromIterator<SystemPathBuf> for SitePackagesPaths {
@ -85,13 +130,67 @@ impl PartialEq<&[SystemPathBuf]> for SitePackagesPaths {
}
#[derive(Debug)]
pub(crate) enum PythonEnvironment {
pub enum PythonEnvironment {
Virtual(VirtualEnvironment),
System(SystemEnvironment),
}
impl PythonEnvironment {
pub(crate) fn new(
pub fn discover(
project_root: &SystemPath,
system: &dyn System,
) -> Result<Option<Self>, SitePackagesDiscoveryError> {
fn resolve_environment(
system: &dyn System,
path: &SystemPath,
origin: SysPrefixPathOrigin,
) -> Result<PythonEnvironment, SitePackagesDiscoveryError> {
tracing::debug!("Resolving {origin}: {path}");
PythonEnvironment::new(path, origin, system)
}
if let Ok(virtual_env) = system.env_var("VIRTUAL_ENV") {
return resolve_environment(
system,
SystemPath::new(&virtual_env),
SysPrefixPathOrigin::VirtualEnvVar,
)
.map(Some);
}
if let Ok(conda_env) = system.env_var("CONDA_PREFIX") {
return resolve_environment(
system,
SystemPath::new(&conda_env),
SysPrefixPathOrigin::CondaPrefixVar,
)
.map(Some);
}
tracing::debug!("Discovering virtual environment in `{project_root}`");
let virtual_env_directory = project_root.join(".venv");
match PythonEnvironment::new(
&virtual_env_directory,
SysPrefixPathOrigin::LocalVenv,
system,
) {
Ok(environment) => return Ok(Some(environment)),
Err(err) => {
if system.is_directory(&virtual_env_directory) {
tracing::debug!(
"Ignoring automatically detected virtual environment at `{}`: {}",
&virtual_env_directory,
err
);
}
}
}
Ok(None)
}
pub fn new(
path: impl AsRef<SystemPath>,
origin: SysPrefixPathOrigin,
system: &dyn System,
@ -111,23 +210,17 @@ impl PythonEnvironment {
}
}
/// Returns the `site-packages` directories for this Python environment,
/// as well as the Python version that was used to create this environment
/// (the latter will only be available for virtual environments that specify
/// Returns the Python version that was used to create this environment
/// (will only be available for virtual environments that specify
/// the metadata in their `pyvenv.cfg` files).
///
/// See the documentation for [`site_packages_directory_from_sys_prefix`] for more details.
pub(crate) fn into_settings(
self,
system: &dyn System,
) -> SitePackagesDiscoveryResult<(SitePackagesPaths, Option<PythonVersionWithSource>)> {
Ok(match self {
Self::Virtual(venv) => (venv.site_packages_directories(system)?, venv.version),
Self::System(env) => (env.site_packages_directories(system)?, None),
})
pub fn python_version_from_metadata(&self) -> Option<&PythonVersionWithSource> {
match self {
Self::Virtual(venv) => venv.version.as_ref(),
Self::System(_) => None,
}
}
fn site_packages_directories(
pub fn site_packages_paths(
&self,
system: &dyn System,
) -> SitePackagesDiscoveryResult<SitePackagesPaths> {
@ -173,7 +266,7 @@ impl PythonImplementation {
/// The format of this file is not defined anywhere, and exactly which keys are present
/// depends on the tool that was used to create the virtual environment.
#[derive(Debug)]
pub(crate) struct VirtualEnvironment {
pub struct VirtualEnvironment {
root_path: SysPrefixPath,
base_executable_home_path: PythonHomePath,
include_system_site_packages: bool,
@ -322,7 +415,7 @@ impl VirtualEnvironment {
);
if let Some(parent_env_site_packages) = parent_environment.as_deref() {
match parent_env_site_packages.site_packages_directories(system) {
match parent_env_site_packages.site_packages_paths(system) {
Ok(parent_environment_site_packages) => {
site_packages_directories.extend(parent_environment_site_packages);
}
@ -492,7 +585,7 @@ struct RawPyvenvCfg<'s> {
/// This environment may or may not be one that is managed by the operating system itself, e.g.,
/// this captures both Homebrew-installed Python versions and the bundled macOS Python installation.
#[derive(Debug)]
pub(crate) struct SystemEnvironment {
pub struct SystemEnvironment {
root_path: SysPrefixPath,
}
@ -1599,7 +1692,7 @@ mod tests {
// directory
let env =
PythonEnvironment::new("/env", SysPrefixPathOrigin::PythonCliFlag, &system).unwrap();
let site_packages = env.site_packages_directories(&system);
let site_packages = env.site_packages_paths(&system);
if cfg!(unix) {
assert!(
matches!(
@ -1638,7 +1731,7 @@ mod tests {
// Environment creation succeeds, but site-packages retrieval fails
let env =
PythonEnvironment::new("/env", SysPrefixPathOrigin::PythonCliFlag, &system).unwrap();
let site_packages = env.site_packages_directories(&system);
let site_packages = env.site_packages_paths(&system);
assert!(
matches!(
site_packages,