mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-08-22 19:34:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

[ty] Implement non-stdlib stub mapping for classes and functions (#19471)

Browse source 
This implements mapping of definitions in stubs to definitions in the
"real" implementation using the approach described in
https://github.com/astral-sh/ty/issues/788#issuecomment-3097000287

I've tested this with goto-definition in vscode with code that uses
`colorama` and `types-colorama`.

Notably this implementation does not add support for stub-mapping stdlib
modules, which can be done as an essentially orthogonal followup in the
implementation of `resolve_real_module`.

Part of https://github.com/astral-sh/ty/issues/788
This commit is contained in:
Aria Desires 2025-07-22 08:42:55 -04:00 • committed by GitHub
parent 6d4687c9af
commit c82fa94e0a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 881 additions and 49 deletions
Download patch file Download diff file Expand all files Collapse all files

4
crates/ty_python_semantic/src/lib.rs
View file

@ -8,7 +8,7 @@ pub use db::Db;
pub use module_name::ModuleName;
pub use module_resolver::{
KnownModule, Module, SearchPathValidationError, SearchPaths, resolve_module,
system_module_search_paths,
resolve_real_module, system_module_search_paths,
};
pub use program::{
Program, ProgramSettings, PythonVersionFileSource, PythonVersionSource,
@ -19,7 +19,7 @@ pub use semantic_model::{Completion, CompletionKind, HasType, NameKind, Semantic
pub use site_packages::{PythonEnvironment, SitePackagesPaths, SysPrefixPathOrigin};
pub use types::ide_support::{
ResolvedDefinition, definitions_for_attribute, definitions_for_imported_symbol,
definitions_for_name,
definitions_for_name, map_stub_definition,
};
pub use util::diagnostics::add_inferred_python_version_hint_to_diagnostic;

2
crates/ty_python_semantic/src/module_resolver/mod.rs
View file

@ -4,7 +4,7 @@ pub use module::{KnownModule, Module};
pub use path::SearchPathValidationError;
pub use resolver::SearchPaths;
pub(crate) use resolver::file_to_module;
pub use resolver::resolve_module;
pub use resolver::{resolve_module, resolve_real_module};
use ruff_db::system::SystemPath;
use crate::Db;

34
crates/ty_python_semantic/src/module_resolver/path.rs
View file

@ -314,7 +314,11 @@ fn query_stdlib_version(
let Some(module_name) = stdlib_path_to_module_name(relative_path) else {
return TypeshedVersionsQueryResult::DoesNotExist;
};
let ResolverContext { db, python_version } = context;
let ResolverContext {
db,
python_version,
mode: _,
} = context;
typeshed_versions(*db).query_module(&module_name, *python_version)
}
@ -701,6 +705,7 @@ mod tests {
use ruff_python_ast::PythonVersion;
use crate::db::tests::TestDb;
use crate::module_resolver::resolver::ModuleResolveMode;
use crate::module_resolver::testing::{FileSpec, MockedTypeshed, TestCase, TestCaseBuilder};
use super::*;
@ -965,7 +970,8 @@ mod tests {
};
let (db, stdlib_path) = py38_typeshed_test_case(TYPESHED);
let resolver = ResolverContext::new(&db, PythonVersion::PY38);
let resolver =
ResolverContext::new(&db, PythonVersion::PY38, ModuleResolveMode::StubsAllowed);
let asyncio_regular_package = stdlib_path.join("asyncio");
assert!(asyncio_regular_package.is_directory(&resolver));
@ -995,7 +1001,8 @@ mod tests {
};
let (db, stdlib_path) = py38_typeshed_test_case(TYPESHED);
let resolver = ResolverContext::new(&db, PythonVersion::PY38);
let resolver =
ResolverContext::new(&db, PythonVersion::PY38, ModuleResolveMode::StubsAllowed);
let xml_namespace_package = stdlib_path.join("xml");
assert!(xml_namespace_package.is_directory(&resolver));
@ -1017,7 +1024,8 @@ mod tests {
};
let (db, stdlib_path) = py38_typeshed_test_case(TYPESHED);
let resolver = ResolverContext::new(&db, PythonVersion::PY38);
let resolver =
ResolverContext::new(&db, PythonVersion::PY38, ModuleResolveMode::StubsAllowed);
let functools_module = stdlib_path.join("functools.pyi");
assert!(functools_module.to_file(&resolver).is_some());
@ -1033,7 +1041,8 @@ mod tests {
};
let (db, stdlib_path) = py38_typeshed_test_case(TYPESHED);
let resolver = ResolverContext::new(&db, PythonVersion::PY38);
let resolver =
ResolverContext::new(&db, PythonVersion::PY38, ModuleResolveMode::StubsAllowed);
let collections_regular_package = stdlib_path.join("collections");
assert_eq!(collections_regular_package.to_file(&resolver), None);
@ -1049,7 +1058,8 @@ mod tests {
};
let (db, stdlib_path) = py38_typeshed_test_case(TYPESHED);
let resolver = ResolverContext::new(&db, PythonVersion::PY38);
let resolver =
ResolverContext::new(&db, PythonVersion::PY38, ModuleResolveMode::StubsAllowed);
let importlib_namespace_package = stdlib_path.join("importlib");
assert_eq!(importlib_namespace_package.to_file(&resolver), None);
@ -1070,7 +1080,8 @@ mod tests {
};
let (db, stdlib_path) = py38_typeshed_test_case(TYPESHED);
let resolver = ResolverContext::new(&db, PythonVersion::PY38);
let resolver =
ResolverContext::new(&db, PythonVersion::PY38, ModuleResolveMode::StubsAllowed);
let non_existent = stdlib_path.join("doesnt_even_exist");
assert_eq!(non_existent.to_file(&resolver), None);
@ -1098,7 +1109,8 @@ mod tests {
};
let (db, stdlib_path) = py39_typeshed_test_case(TYPESHED);
let resolver = ResolverContext::new(&db, PythonVersion::PY39);
let resolver =
ResolverContext::new(&db, PythonVersion::PY39, ModuleResolveMode::StubsAllowed);
// Since we've set the target version to Py39,
// `collections` should now exist as a directory, according to VERSIONS...
@ -1129,7 +1141,8 @@ mod tests {
};
let (db, stdlib_path) = py39_typeshed_test_case(TYPESHED);
let resolver = ResolverContext::new(&db, PythonVersion::PY39);
let resolver =
ResolverContext::new(&db, PythonVersion::PY39, ModuleResolveMode::StubsAllowed);
// The `importlib` directory now also exists
let importlib_namespace_package = stdlib_path.join("importlib");
@ -1153,7 +1166,8 @@ mod tests {
};
let (db, stdlib_path) = py39_typeshed_test_case(TYPESHED);
let resolver = ResolverContext::new(&db, PythonVersion::PY39);
let resolver =
ResolverContext::new(&db, PythonVersion::PY39, ModuleResolveMode::StubsAllowed);
// The `xml` package no longer exists on py39:
let xml_namespace_package = stdlib_path.join("xml");

66
crates/ty_python_semantic/src/module_resolver/resolver.rs
View file

@ -21,11 +21,32 @@ use super::path::{ModulePath, SearchPath, SearchPathValidationError, SystemOrVen
/// Resolves a module name to a module.
pub fn resolve_module(db: &dyn Db, module_name: &ModuleName) -> Option<Module> {
let interned_name = ModuleNameIngredient::new(db, module_name);
let interned_name = ModuleNameIngredient::new(db, module_name, ModuleResolveMode::StubsAllowed);
resolve_module_query(db, interned_name)
}
/// Resolves a module name to a module (stubs not allowed).
pub fn resolve_real_module(db: &dyn Db, module_name: &ModuleName) -> Option<Module> {
let interned_name =
ModuleNameIngredient::new(db, module_name, ModuleResolveMode::StubsNotAllowed);
resolve_module_query(db, interned_name)
}
/// Which files should be visible when doing a module query
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(crate) enum ModuleResolveMode {
StubsAllowed,
StubsNotAllowed,
}
impl ModuleResolveMode {
fn stubs_allowed(self) -> bool {
matches!(self, Self::StubsAllowed)
}
}
/// Salsa query that resolves an interned [`ModuleNameIngredient`] to a module.
///
/// This query should not be called directly. Instead, use [`resolve_module`]. It only exists
@ -36,9 +57,10 @@ pub(crate) fn resolve_module_query<'db>(
module_name: ModuleNameIngredient<'db>,
) -> Option<Module> {
let name = module_name.name(db);
let mode = module_name.mode(db);
let _span = tracing::trace_span!("resolve_module", %name).entered();
let Some(resolved) = resolve_name(db, name) else {
let Some(resolved) = resolve_name(db, name, mode) else {
tracing::debug!("Module `{name}` not found in search paths");
return None;
};
@ -514,6 +536,7 @@ impl<'db> Iterator for PthFileIterator<'db> {
struct ModuleNameIngredient<'db> {
#[returns(ref)]
pub(super) name: ModuleName,
pub(super) mode: ModuleResolveMode,
}
/// Returns `true` if the module name refers to a standard library module which can't be shadowed
@ -528,10 +551,10 @@ fn is_non_shadowable(minor_version: u8, module_name: &str) -> bool {
/// Given a module name and a list of search paths in which to lookup modules,
/// attempt to resolve the module name
fn resolve_name(db: &dyn Db, name: &ModuleName) -> Option<ResolvedName> {
fn resolve_name(db: &dyn Db, name: &ModuleName, mode: ModuleResolveMode) -> Option<ResolvedName> {
let program = Program::get(db);
let python_version = program.python_version(db);
let resolver_state = ResolverContext::new(db, python_version);
let resolver_state = ResolverContext::new(db, python_version, mode);
let is_non_shadowable = is_non_shadowable(python_version.minor, name.as_str());
let name = RelaxedModuleName::new(name);
@ -548,7 +571,7 @@ fn resolve_name(db: &dyn Db, name: &ModuleName) -> Option<ResolvedName> {
continue;
}
if !search_path.is_standard_library() {
if !search_path.is_standard_library() && resolver_state.mode.stubs_allowed() {
match resolve_name_in_search_path(&resolver_state, &stub_name, search_path) {
Ok((package_kind, ResolvedName::FileModule(module))) => {
if package_kind.is_root() && module.kind.is_module() {
@ -717,14 +740,16 @@ fn resolve_name_in_search_path(
/// resolving modules.
fn resolve_file_module(module: &ModulePath, resolver_state: &ResolverContext) -> Option<File> {
// Stubs have precedence over source files
let file = module
.with_pyi_extension()
.to_file(resolver_state)
.or_else(|| {
module
.with_py_extension()
.and_then(|path| path.to_file(resolver_state))
})?;
let stub_file = if resolver_state.mode.stubs_allowed() {
module.with_pyi_extension().to_file(resolver_state)
} else {
None
};
let file = stub_file.or_else(|| {
module
.with_py_extension()
.and_then(|path| path.to_file(resolver_state))
})?;
// For system files, test if the path has the correct casing.
// We can skip this step for vendored files or virtual files because
@ -833,11 +858,20 @@ impl PackageKind {
pub(super) struct ResolverContext<'db> {
pub(super) db: &'db dyn Db,
pub(super) python_version: PythonVersion,
pub(super) mode: ModuleResolveMode,
}
impl<'db> ResolverContext<'db> {
pub(super) fn new(db: &'db dyn Db, python_version: PythonVersion) -> Self {
Self { db, python_version }
pub(super) fn new(
db: &'db dyn Db,
python_version: PythonVersion,
mode: ModuleResolveMode,
) -> Self {
Self {
db,
python_version,
mode,
}
}
pub(super) fn vendored(&self) -> &VendoredFileSystem {
@ -1539,7 +1573,7 @@ mod tests {
assert_function_query_was_not_run(
&db,
resolve_module_query,
ModuleNameIngredient::new(&db, functools_module_name),
ModuleNameIngredient::new(&db, functools_module_name, ModuleResolveMode::StubsAllowed),
&events,
);
assert_eq!(functools_module.search_path().unwrap(), &stdlib);

241
crates/ty_python_semantic/src/types/ide_support.rs
View file

@ -20,7 +20,7 @@ use ruff_python_ast::name::Name;
use ruff_text_size::{Ranged, TextRange};
use rustc_hash::FxHashSet;
pub use resolve_definition::ResolvedDefinition;
pub use resolve_definition::{ResolvedDefinition, map_stub_definition};
use resolve_definition::{find_symbol_in_scope, resolve_definition};
pub(crate) fn all_declarations_and_bindings<'db>(
@ -788,16 +788,19 @@ mod resolve_definition {
//! "resolved definitions". This is done recursively to find the original
//! definition targeted by the import.
use indexmap::IndexSet;
use ruff_db::files::{File, FileRange};
use ruff_db::parsed::parsed_module;
use ruff_db::parsed::{ParsedModuleRef, parsed_module};
use ruff_python_ast as ast;
use ruff_text_size::TextRange;
use ruff_text_size::{Ranged, TextRange};
use rustc_hash::FxHashSet;
use tracing::trace;
use crate::module_resolver::file_to_module;
use crate::semantic_index::definition::{Definition, DefinitionKind};
use crate::semantic_index::place::ScopeId;
use crate::semantic_index::{global_scope, place_table, use_def_map};
use crate::{Db, ModuleName, resolve_module};
use crate::semantic_index::place::{NodeWithScopeKind, ScopeId};
use crate::semantic_index::{global_scope, place_table, semantic_index, use_def_map};
use crate::{Db, ModuleName, resolve_module, resolve_real_module};
/// Represents the result of resolving an import to either a specific definition or
/// a specific range within a file.
@ -812,6 +815,15 @@ mod resolve_definition {
FileWithRange(FileRange),
}
impl<'db> ResolvedDefinition<'db> {
fn file(&self, db: &'db dyn Db) -> File {
match self {
ResolvedDefinition::Definition(definition) => definition.file(db),
ResolvedDefinition::FileWithRange(file_range) => file_range.file(),
}
}
}
/// Resolve import definitions to their targets.
/// Returns resolved definitions which can be either specific definitions or module files.
/// For non-import definitions, returns the definition wrapped in `ResolvedDefinition::Definition`.
@ -954,14 +966,14 @@ mod resolve_definition {
db: &'db dyn Db,
scope: ScopeId<'db>,
symbol_name: &str,
) -> Vec<Definition<'db>> {
) -> IndexSet<Definition<'db>> {
let place_table = place_table(db, scope);
let Some(place_id) = place_table.place_id_by_name(symbol_name) else {
return Vec::new();
return IndexSet::new();
};
let use_def_map = use_def_map(db, scope);
let mut definitions = Vec::new();
let mut definitions = IndexSet::new();
// Get all definitions (both bindings and declarations) for this place
let bindings = use_def_map.all_reachable_bindings(place_id);
@ -969,16 +981,223 @@ mod resolve_definition {
for binding in bindings {
if let Some(def) = binding.binding.definition() {
definitions.push(def);
definitions.insert(def);
}
}
for declaration in declarations {
if let Some(def) = declaration.declaration.definition() {
definitions.push(def);
definitions.insert(def);
}
}
definitions
}
/// Given a definition that may be in a stub file, find the "real" definition in a non-stub.
#[tracing::instrument(skip_all)]
pub fn map_stub_definition<'db>(
db: &'db dyn Db,
def: &ResolvedDefinition<'db>,
) -> Option<Vec<ResolvedDefinition<'db>>> {
trace!("Stub mapping definition...");
// If the file isn't a stub, this is presumably the real definition
let stub_file = def.file(db);
if !stub_file.is_stub(db) {
trace!("File isn't a stub, no stub mapping to do");
return None;
}
// It's definitely a stub, so now rerun module resolution but with stubs disabled.
let stub_module = file_to_module(db, stub_file)?;
trace!("Found stub module: {}", stub_module.name());
let real_module = resolve_real_module(db, stub_module.name())?;
trace!("Found real module: {}", real_module.name());
let real_file = real_module.file()?;
trace!("Found real file: {}", real_file.path(db));
// A definition has a "Definition Path" in a file made of nested definitions (~scopes):
//
// ```
// class myclass: # ./myclass
// def some_func(args: bool): # ./myclass/some_func
// # ^~~~ ./myclass/other_func/args/
// ```
//
// So our heuristic goal here is to compute a Definition Path in the stub file
// and then resolve the same Definition Path in the real file.
//
// NOTE: currently a path component is just a str, but in the future additional
// disambiguators (like "is a class def") could be added if needed.
let mut path = Vec::new();
let stub_parsed;
let stub_ref;
match *def {
ResolvedDefinition::Definition(definition) => {
stub_parsed = parsed_module(db, stub_file);
stub_ref = stub_parsed.load(db);
// Get the leaf of the path (the definition itself)
let leaf = definition_path_component_for_leaf(db, &stub_ref, definition)
.map_err(|()| {
trace!("Found unsupported DefinitionKind while stub mapping, giving up");
})
.ok()?;
path.push(leaf);
// Get the ancestors of the path (all the definitions we're nested under)
let index = semantic_index(db, stub_file);
for (_scope_id, scope) in index.ancestor_scopes(definition.file_scope(db)) {
let node = scope.node();
let component = definition_path_component_for_node(&stub_ref, node)
.map_err(|()| {
trace!("Found unsupported NodeScopeKind while stub mapping, giving up");
})
.ok()?;
if let Some(component) = component {
path.push(component);
}
}
trace!("Built Definition Path: {path:?}");
}
ResolvedDefinition::FileWithRange(file_range) => {
return if file_range.range() == TextRange::default() {
trace!(
"Found module mapping: {} => {}",
stub_file.path(db),
real_file.path(db)
);
// This is just a reference to a module, no need to do paths
Some(vec![ResolvedDefinition::FileWithRange(FileRange::new(
real_file,
TextRange::default(),
))])
} else {
// Not yet implemented -- in this case we want to recover something like a Definition
// and build a Definition Path, but this input is a bit too abstract for now.
trace!("Found arbitrary FileWithRange by stub mapping, giving up");
None
};
}
}
// Walk down the Definition Path in the real file
let mut definitions = Vec::new();
let index = semantic_index(db, real_file);
let real_parsed = parsed_module(db, real_file);
let real_ref = real_parsed.load(db);
// Start our search in the module (global) scope
let mut scopes = vec![global_scope(db, real_file)];
while let Some(component) = path.pop() {
trace!("Traversing definition path component: {}", component);
// We're doing essentially a breadth-first traversal of the definitions.
// If ever we find multiple matching scopes for a component, we need to continue
// walking down each of them to try to resolve the path. Here we loop over
// all the scopes at the current level of search.
for scope in std::mem::take(&mut scopes) {
if path.is_empty() {
// We're at the end of the path, everything we find here is the final result
definitions.extend(
find_symbol_in_scope(db, scope, component)
.into_iter()
.map(ResolvedDefinition::Definition),
);
} else {
// We're in the middle of the path, look for scopes that match the current component
for (child_scope_id, child_scope) in index.child_scopes(scope.file_scope_id(db))
{
let scope_node = child_scope.node();
if let Ok(Some(real_component)) =
definition_path_component_for_node(&real_ref, scope_node)
{
if real_component == component {
scopes.push(child_scope_id.to_scope_id(db, real_file));
}
}
scope.node(db);
}
}
}
trace!(
"Found {} scopes and {} definitions",
scopes.len(),
definitions.len()
);
}
if definitions.is_empty() {
trace!("No definitions found in real file, stub mapping failed");
None
} else {
trace!("Found {} definitions from stub mapping", definitions.len());
Some(definitions)
}
}
/// Computes a "Definition Path" component for an internal node of the definition path.
///
/// See [`map_stub_definition`][] for details.
fn definition_path_component_for_node<'parse>(
parsed: &'parse ParsedModuleRef,
node: &NodeWithScopeKind,
) -> Result<Option<&'parse str>, ()> {
let component = match node {
NodeWithScopeKind::Module => {
// This is just implicit, so has no component
return Ok(None);
}
NodeWithScopeKind::Class(class) => class.node(parsed).name.as_str(),
NodeWithScopeKind::Function(func) => func.node(parsed).name.as_str(),
NodeWithScopeKind::TypeAlias(_)
| NodeWithScopeKind::ClassTypeParameters(_)
| NodeWithScopeKind::FunctionTypeParameters(_)
| NodeWithScopeKind::TypeAliasTypeParameters(_)
| NodeWithScopeKind::Lambda(_)
| NodeWithScopeKind::ListComprehension(_)
| NodeWithScopeKind::SetComprehension(_)
| NodeWithScopeKind::DictComprehension(_)
| NodeWithScopeKind::GeneratorExpression(_) => {
// Not yet implemented
return Err(());
}
};
Ok(Some(component))
}
/// Computes a "Definition Path" component for a leaf node of the definition path.
///
/// See [`map_stub_definition`][] for details.
fn definition_path_component_for_leaf<'parse>(
db: &dyn Db,
parsed: &'parse ParsedModuleRef,
definition: Definition,
) -> Result<&'parse str, ()> {
let component = match definition.kind(db) {
DefinitionKind::Function(func) => func.node(parsed).name.as_str(),
DefinitionKind::Class(class) => class.node(parsed).name.as_str(),
DefinitionKind::TypeAlias(_)
| DefinitionKind::Import(_)
| DefinitionKind::ImportFrom(_)
| DefinitionKind::StarImport(_)
| DefinitionKind::NamedExpression(_)
| DefinitionKind::Assignment(_)
| DefinitionKind::AnnotatedAssignment(_)
| DefinitionKind::AugmentedAssignment(_)
| DefinitionKind::For(_)
| DefinitionKind::Comprehension(_)
| DefinitionKind::VariadicPositionalParameter(_)
| DefinitionKind::VariadicKeywordParameter(_)
| DefinitionKind::Parameter(_)
| DefinitionKind::WithItem(_)
| DefinitionKind::MatchPattern(_)
| DefinitionKind::ExceptHandler(_)
| DefinitionKind::TypeVar(_)
| DefinitionKind::ParamSpec(_)
| DefinitionKind::TypeVarTuple(_) => {
// Not yet implemented
return Err(());
}
};
Ok(component)
}
}