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

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

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",

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 }

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 mymo<CURSOR>dule 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 = mymo<CURSOR>dule
+",
+            )
+            .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_func<CURSOR>tion())
+",
+            )
+            .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_func<CURSOR>tion())
+",
+            )
+            .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 = MyC<CURSOR>lass
+",
+            )
+            .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 = MyCl<CURSOR>ass(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.act<CURSOR>ion()
+",
+            )
+            .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.act<CURSOR>ion()
+",
+            )
+            .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 MyC<CURSOR>lass
+",
+            )
+            .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
+        }
+    }
+}

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<ResolvedDefinition<'db>> {
+    ) -> impl Iterator<Item = ResolvedDefinition<'db>> {
-        // TODO: Implement stub-to-source mapping logic
+        if let Some(definitions) = map_stub_definition(self.db, &def) {
-        // For now, just return the original definition
+            return Either::Left(definitions.into_iter());
-        vec![def]
+        }
+        Either::Right(std::iter::once(def))
     }
 
     /// Map multiple `ResolvedDefinitions`, applying stub-to-source mapping to each.

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;
 

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;

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");

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<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,10 +740,12 @@ 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
+    let stub_file = if resolver_state.mode.stubs_allowed() {
-        .with_pyi_extension()
+        module.with_pyi_extension().to_file(resolver_state)
-        .to_file(resolver_state)
+    } else {
-        .or_else(|| {
+        None
+    };
+    let file = stub_file.or_else(|| {
         module
             .with_py_extension()
             .and_then(|path| path.to_file(resolver_state))
@@ -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 {
+    pub(super) fn new(
-        Self { db, python_version }
+        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);

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::place::{NodeWithScopeKind, ScopeId};
-    use crate::semantic_index::{global_scope, place_table, use_def_map};
+    use crate::semantic_index::{global_scope, place_table, semantic_index, use_def_map};
-    use crate::{Db, ModuleName, resolve_module};
+    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)
+    }
 }