diff --git a/Cargo.lock b/Cargo.lock index 4911afc3d9..b0d761cfee 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4211,6 +4211,7 @@ version = "0.0.0" dependencies = [ "bitflags 2.9.1", "insta", + "itertools 0.14.0", "regex", "ruff_db", "ruff_python_ast", diff --git a/crates/ty_ide/Cargo.toml b/crates/ty_ide/Cargo.toml index f54aa96ae7..13efcda3ca 100644 --- a/crates/ty_ide/Cargo.toml +++ b/crates/ty_ide/Cargo.toml @@ -20,6 +20,7 @@ ruff_source_file = { workspace = true } ruff_text_size = { workspace = true } ty_python_semantic = { workspace = true } +itertools = { workspace = true } regex = { workspace = true } rustc-hash = { workspace = true } salsa = { workspace = true } diff --git a/crates/ty_ide/src/goto_definition.rs b/crates/ty_ide/src/goto_definition.rs index 3a65d9fc3f..4c4d090fb3 100644 --- a/crates/ty_ide/src/goto_definition.rs +++ b/crates/ty_ide/src/goto_definition.rs @@ -29,3 +29,568 @@ pub fn goto_definition( value: definition_targets, }) } + +#[cfg(test)] +mod test { + use crate::tests::{CursorTest, IntoDiagnostic}; + use crate::{NavigationTarget, goto_definition}; + use insta::assert_snapshot; + use ruff_db::diagnostic::{ + Annotation, Diagnostic, DiagnosticId, LintName, Severity, Span, SubDiagnostic, + }; + use ruff_db::files::FileRange; + use ruff_text_size::Ranged; + + /// goto-definition on a module should go to the .py not the .pyi + /// + /// TODO: this currently doesn't work right! This is especially surprising + /// because [`goto_definition_stub_map_module_ref`] works fine. + #[test] + fn goto_definition_stub_map_module_import() { + let test = CursorTest::builder() + .source( + "main.py", + " +from mymodule import my_function +", + ) + .source( + "mymodule.py", + r#" +def my_function(): + return "hello" +"#, + ) + .source( + "mymodule.pyi", + r#" +def my_function(): ... +"#, + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> mymodule.pyi:1:1 + | + 1 | + | ^ + 2 | def my_function(): ... + | + info: Source + --> main.py:2:6 + | + 2 | from mymodule import my_function + | ^^^^^^^^ + | + "); + } + + /// goto-definition on a module ref should go to the .py not the .pyi + #[test] + fn goto_definition_stub_map_module_ref() { + let test = CursorTest::builder() + .source( + "main.py", + " +import mymodule +x = mymodule +", + ) + .source( + "mymodule.py", + r#" +def my_function(): + return "hello" +"#, + ) + .source( + "mymodule.pyi", + r#" +def my_function(): ... +"#, + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r#" + info[goto-definition]: Definition + --> mymodule.py:1:1 + | + 1 | + | ^ + 2 | def my_function(): + 3 | return "hello" + | + info: Source + --> main.py:3:5 + | + 2 | import mymodule + 3 | x = mymodule + | ^^^^^^^^ + | + "#); + } + + /// goto-definition on a function call should go to the .py not the .pyi + #[test] + fn goto_definition_stub_map_function() { + let test = CursorTest::builder() + .source( + "main.py", + " +from mymodule import my_function +print(my_function()) +", + ) + .source( + "mymodule.py", + r#" +def my_function(): + return "hello" + +def other_function(): + return "other" +"#, + ) + .source( + "mymodule.pyi", + r#" +def my_function(): ... + +def other_function(): ... +"#, + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r#" + info[goto-definition]: Definition + --> mymodule.py:2:5 + | + 2 | def my_function(): + | ^^^^^^^^^^^ + 3 | return "hello" + | + info: Source + --> main.py:3:7 + | + 2 | from mymodule import my_function + 3 | print(my_function()) + | ^^^^^^^^^^^ + | + "#); + } + + /// goto-definition on a function that's redefined many times in the impl .py + /// + /// Currently this yields all instances. There's an argument for only yielding + /// the final one since that's the one "exported" but, this is consistent for + /// how we do file-local goto-definition. + #[test] + fn goto_definition_stub_map_function_redefine() { + let test = CursorTest::builder() + .source( + "main.py", + " +from mymodule import my_function +print(my_function()) +", + ) + .source( + "mymodule.py", + r#" +def my_function(): + return "hello" + +def my_function(): + return "hello again" + +def my_function(): + return "we can't keep doing this" + +def other_function(): + return "other" +"#, + ) + .source( + "mymodule.pyi", + r#" +def my_function(): ... + +def other_function(): ... +"#, + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r#" + info[goto-definition]: Definition + --> mymodule.py:2:5 + | + 2 | def my_function(): + | ^^^^^^^^^^^ + 3 | return "hello" + | + info: Source + --> main.py:3:7 + | + 2 | from mymodule import my_function + 3 | print(my_function()) + | ^^^^^^^^^^^ + | + + info[goto-definition]: Definition + --> mymodule.py:5:5 + | + 3 | return "hello" + 4 | + 5 | def my_function(): + | ^^^^^^^^^^^ + 6 | return "hello again" + | + info: Source + --> main.py:3:7 + | + 2 | from mymodule import my_function + 3 | print(my_function()) + | ^^^^^^^^^^^ + | + + info[goto-definition]: Definition + --> mymodule.py:8:5 + | + 6 | return "hello again" + 7 | + 8 | def my_function(): + | ^^^^^^^^^^^ + 9 | return "we can't keep doing this" + | + info: Source + --> main.py:3:7 + | + 2 | from mymodule import my_function + 3 | print(my_function()) + | ^^^^^^^^^^^ + | + "#); + } + + /// goto-definition on a class ref go to the .py not the .pyi + #[test] + fn goto_definition_stub_map_class_ref() { + let test = CursorTest::builder() + .source( + "main.py", + " +from mymodule import MyClass +x = MyClass +", + ) + .source( + "mymodule.py", + r#" +class MyClass: + def __init__(self, val): + self.val = val + +class MyOtherClass: + def __init__(self, val): + self.val = val + 1 +"#, + ) + .source( + "mymodule.pyi", + r#" +class MyClass: + def __init__(self, val: bool): ... + +class MyOtherClass: + def __init__(self, val: bool): ... +"#, + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> mymodule.py:2:7 + | + 2 | class MyClass: + | ^^^^^^^ + 3 | def __init__(self, val): + 4 | self.val = val + | + info: Source + --> main.py:3:5 + | + 2 | from mymodule import MyClass + 3 | x = MyClass + | ^^^^^^^ + | + "); + } + + /// goto-definition on a class init should go to the .py not the .pyi + #[test] + fn goto_definition_stub_map_class_init() { + let test = CursorTest::builder() + .source( + "main.py", + " +from mymodule import MyClass +x = MyClass(0) +", + ) + .source( + "mymodule.py", + r#" +class MyClass: + def __init__(self, val): + self.val = val + +class MyOtherClass: + def __init__(self, val): + self.val = val + 1 +"#, + ) + .source( + "mymodule.pyi", + r#" +class MyClass: + def __init__(self, val: bool): ... + +class MyOtherClass: + def __init__(self, val: bool): ... +"#, + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> mymodule.py:2:7 + | + 2 | class MyClass: + | ^^^^^^^ + 3 | def __init__(self, val): + 4 | self.val = val + | + info: Source + --> main.py:3:5 + | + 2 | from mymodule import MyClass + 3 | x = MyClass(0) + | ^^^^^^^ + | + "); + } + + /// goto-definition on a class method should go to the .py not the .pyi + #[test] + fn goto_definition_stub_map_class_method() { + let test = CursorTest::builder() + .source( + "main.py", + " +from mymodule import MyClass +x = MyClass(0) +x.action() +", + ) + .source( + "mymodule.py", + r#" +class MyClass: + def __init__(self, val): + self.val = val + def action(self): + print(self.val) + +class MyOtherClass: + def __init__(self, val): + self.val = val + 1 +"#, + ) + .source( + "mymodule.pyi", + r#" +class MyClass: + def __init__(self, val: bool): ... + def action(self): ... + +class MyOtherClass: + def __init__(self, val: bool): ... +"#, + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> mymodule.py:5:9 + | + 3 | def __init__(self, val): + 4 | self.val = val + 5 | def action(self): + | ^^^^^^ + 6 | print(self.val) + | + info: Source + --> main.py:4:1 + | + 2 | from mymodule import MyClass + 3 | x = MyClass(0) + 4 | x.action() + | ^^^^^^^^ + | + "); + } + + /// goto-definition on a class function should go to the .py not the .pyi + #[test] + fn goto_definition_stub_map_class_function() { + let test = CursorTest::builder() + .source( + "main.py", + " +from mymodule import MyClass +x = MyClass.action() +", + ) + .source( + "mymodule.py", + r#" +class MyClass: + def __init__(self, val): + self.val = val + def action(): + print("hi!") + +class MyOtherClass: + def __init__(self, val): + self.val = val + 1 +"#, + ) + .source( + "mymodule.pyi", + r#" +class MyClass: + def __init__(self, val: bool): ... + def action(): ... + +class MyOtherClass: + def __init__(self, val: bool): ... +"#, + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r#" + info[goto-definition]: Definition + --> mymodule.py:5:9 + | + 3 | def __init__(self, val): + 4 | self.val = val + 5 | def action(): + | ^^^^^^ + 6 | print("hi!") + | + info: Source + --> main.py:3:5 + | + 2 | from mymodule import MyClass + 3 | x = MyClass.action() + | ^^^^^^^^^^^^^^ + | + "#); + } + + /// goto-definition on a class import should go to the .py not the .pyi + #[test] + fn goto_definition_stub_map_class_import() { + let test = CursorTest::builder() + .source( + "main.py", + " +from mymodule import MyClass +", + ) + .source( + "mymodule.py", + r#" +class MyClass: ... +"#, + ) + .source( + "mymodule.pyi", + r#" +class MyClass: ... +"#, + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> mymodule.py:2:7 + | + 2 | class MyClass: ... + | ^^^^^^^ + | + info: Source + --> main.py:2:22 + | + 2 | from mymodule import MyClass + | ^^^^^^^ + | + "); + } + + impl CursorTest { + fn goto_definition(&self) -> String { + let Some(targets) = goto_definition(&self.db, self.cursor.file, self.cursor.offset) + else { + return "No goto target found".to_string(); + }; + + if targets.is_empty() { + return "No definitions found".to_string(); + } + + let source = targets.range; + self.render_diagnostics( + targets + .into_iter() + .map(|target| GotoDefinitionDiagnostic::new(source, &target)), + ) + } + } + + struct GotoDefinitionDiagnostic { + source: FileRange, + target: FileRange, + } + + impl GotoDefinitionDiagnostic { + fn new(source: FileRange, target: &NavigationTarget) -> Self { + Self { + source, + target: FileRange::new(target.file(), target.focus_range()), + } + } + } + + impl IntoDiagnostic for GotoDefinitionDiagnostic { + fn into_diagnostic(self) -> Diagnostic { + let mut source = SubDiagnostic::new(Severity::Info, "Source"); + source.annotate(Annotation::primary( + Span::from(self.source.file()).with_range(self.source.range()), + )); + + let mut main = Diagnostic::new( + DiagnosticId::Lint(LintName::of("goto-definition")), + Severity::Info, + "Definition".to_string(), + ); + main.annotate(Annotation::primary( + Span::from(self.target.file()).with_range(self.target.range()), + )); + main.sub(source); + + main + } + } +} diff --git a/crates/ty_ide/src/stub_mapping.rs b/crates/ty_ide/src/stub_mapping.rs index 9be4ccc7d5..d28823d145 100644 --- a/crates/ty_ide/src/stub_mapping.rs +++ b/crates/ty_ide/src/stub_mapping.rs @@ -1,4 +1,5 @@ -use ty_python_semantic::ResolvedDefinition; +use itertools::Either; +use ty_python_semantic::{ResolvedDefinition, map_stub_definition}; /// Maps `ResolvedDefinitions` from stub files to corresponding definitions in source files. /// @@ -7,12 +8,10 @@ use ty_python_semantic::ResolvedDefinition; /// other language server providers (like hover, completion, and signature help) to find /// docstrings for functions that resolve to stubs. pub(crate) struct StubMapper<'db> { - #[allow(dead_code)] // Will be used when implementation is added db: &'db dyn crate::Db, } impl<'db> StubMapper<'db> { - #[allow(dead_code)] // Will be used in the future pub(crate) fn new(db: &'db dyn crate::Db) -> Self { Self { db } } @@ -21,15 +20,14 @@ impl<'db> StubMapper<'db> { /// /// If the definition is in a stub file and a corresponding source file definition exists, /// returns the source file definition(s). Otherwise, returns the original definition. - #[allow(dead_code)] // Will be used when implementation is added - #[allow(clippy::unused_self)] // Will use self when implementation is added pub(crate) fn map_definition( &self, def: ResolvedDefinition<'db>, - ) -> Vec> { - // TODO: Implement stub-to-source mapping logic - // For now, just return the original definition - vec![def] + ) -> impl Iterator> { + if let Some(definitions) = map_stub_definition(self.db, &def) { + return Either::Left(definitions.into_iter()); + } + Either::Right(std::iter::once(def)) } /// Map multiple `ResolvedDefinitions`, applying stub-to-source mapping to each. diff --git a/crates/ty_python_semantic/src/lib.rs b/crates/ty_python_semantic/src/lib.rs index 57cf070dab..1b8314f43b 100644 --- a/crates/ty_python_semantic/src/lib.rs +++ b/crates/ty_python_semantic/src/lib.rs @@ -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; diff --git a/crates/ty_python_semantic/src/module_resolver/mod.rs b/crates/ty_python_semantic/src/module_resolver/mod.rs index b041ea14ab..2c6548570f 100644 --- a/crates/ty_python_semantic/src/module_resolver/mod.rs +++ b/crates/ty_python_semantic/src/module_resolver/mod.rs @@ -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; diff --git a/crates/ty_python_semantic/src/module_resolver/path.rs b/crates/ty_python_semantic/src/module_resolver/path.rs index 92113426c7..91488d2442 100644 --- a/crates/ty_python_semantic/src/module_resolver/path.rs +++ b/crates/ty_python_semantic/src/module_resolver/path.rs @@ -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"); diff --git a/crates/ty_python_semantic/src/module_resolver/resolver.rs b/crates/ty_python_semantic/src/module_resolver/resolver.rs index 5c670e6d2d..2c576fa998 100644 --- a/crates/ty_python_semantic/src/module_resolver/resolver.rs +++ b/crates/ty_python_semantic/src/module_resolver/resolver.rs @@ -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 { - 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 { + 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 { 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 { +fn resolve_name(db: &dyn Db, name: &ModuleName, mode: ModuleResolveMode) -> Option { 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 { 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 { // 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); diff --git a/crates/ty_python_semantic/src/types/ide_support.rs b/crates/ty_python_semantic/src/types/ide_support.rs index b0eb2dfb93..2ab52ef0fa 100644 --- a/crates/ty_python_semantic/src/types/ide_support.rs +++ b/crates/ty_python_semantic/src/types/ide_support.rs @@ -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> { + ) -> IndexSet> { 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>> { + 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, ()> { + 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) + } }