[ty] Support goto-definition on vendored typeshed stubs (#21020)

This is an alternative to #21012 that more narrowly handles this logic
in the stub-mapping machinery rather than pervasively allowing us to
identify cached files as typeshed stubs. Much of the logic is the same
(pulling the logic out of ty_server so it can be reused).

I don't have a good sense for if one approach is "better" or "worse" in
terms of like, semantics and Weird Bugs that this can cause. This one is
just "less spooky in its broad consequences" and "less muddying of
separation of concerns" and puts the extra logic on a much colder path.
I won't be surprised if one day the previous implementation needs to be
revisited for its more sweeping effects but for now this is good.

Fixes https://github.com/astral-sh/ty/issues/1054

Cargo.lock

@@ -4376,6 +4376,7 @@ dependencies = [
  "tracing",
  "ty_project",
  "ty_python_semantic",
+ "ty_vendored",
 ]
 
 [[package]]
@@ -4504,7 +4505,6 @@ dependencies = [
  "ty_ide",
  "ty_project",
  "ty_python_semantic",
- "ty_vendored",
 ]
 
 [[package]]

crates/ty_ide/Cargo.toml

@@ -25,6 +25,7 @@ ruff_source_file = { workspace = true }
 ruff_text_size = { workspace = true }
 ty_python_semantic = { workspace = true }
 ty_project = { workspace = true, features = ["testing"] }
+ty_vendored = { workspace = true }
 
 get-size2 = { workspace = true }
 itertools = { workspace = true }

crates/ty_ide/src/lib.rs

@@ -45,7 +45,11 @@ pub use signature_help::{ParameterDetails, SignatureDetails, SignatureHelpInfo,
 pub use symbols::{FlatSymbols, HierarchicalSymbols, SymbolId, SymbolInfo, SymbolKind};
 pub use workspace_symbols::{WorkspaceSymbolInfo, workspace_symbols};
 
-use ruff_db::files::{File, FileRange};
+use ruff_db::{
+    files::{File, FileRange},
+    system::SystemPathBuf,
+    vendored::VendoredPath,
+};
 use ruff_text_size::{Ranged, TextRange};
 use rustc_hash::FxHashSet;
 use std::ops::{Deref, DerefMut};
@@ -287,6 +291,38 @@ impl HasNavigationTargets for TypeDefinition<'_> {
     }
 }
 
+/// Get the cache-relative path where vendored paths should be written to.
+pub fn relative_cached_vendored_root() -> SystemPathBuf {
+    // The vendored files are uniquely identified by the source commit.
+    SystemPathBuf::from(format!("vendored/typeshed/{}", ty_vendored::SOURCE_COMMIT))
+}
+
+/// Get the cached version of a vendored path in the cache, ensuring the file is written to disk.
+pub fn cached_vendored_path(
+    db: &dyn ty_python_semantic::Db,
+    path: &VendoredPath,
+) -> Option<SystemPathBuf> {
+    let writable = db.system().as_writable()?;
+    let mut relative_path = relative_cached_vendored_root();
+    relative_path.push(path.as_str());
+
+    // Extract the vendored file onto the system.
+    writable
+        .get_or_cache(&relative_path, &|| db.vendored().read_to_string(path))
+        .ok()
+        .flatten()
+}
+
+/// Get the absolute root path of all cached vendored paths.
+///
+/// This does not ensure that this path exists (this is only used for mapping cached paths
+/// back to vendored ones, so this only matters if we've already been handed a path inside here).
+pub fn cached_vendored_root(db: &dyn ty_python_semantic::Db) -> Option<SystemPathBuf> {
+    let writable = db.system().as_writable()?;
+    let relative_root = relative_cached_vendored_root();
+    Some(writable.cache_dir()?.join(relative_root))
+}
+
 #[cfg(test)]
 mod tests {
     use camino::Utf8Component;

crates/ty_ide/src/stub_mapping.rs

@@ -1,6 +1,9 @@
 use itertools::Either;
+use ruff_db::system::SystemPathBuf;
 use ty_python_semantic::{ResolvedDefinition, map_stub_definition};
 
+use crate::cached_vendored_root;
+
 /// Maps `ResolvedDefinitions` from stub files to corresponding definitions in source files.
 ///
 /// This mapper is used to implement "Go To Definition" functionality that navigates from
@@ -9,11 +12,16 @@ use ty_python_semantic::{ResolvedDefinition, map_stub_definition};
 /// docstrings for functions that resolve to stubs.
 pub(crate) struct StubMapper<'db> {
     db: &'db dyn crate::Db,
+    cached_vendored_root: Option<SystemPathBuf>,
 }
 
 impl<'db> StubMapper<'db> {
     pub(crate) fn new(db: &'db dyn crate::Db) -> Self {
-        Self { db }
+        let cached_vendored_root = cached_vendored_root(db);
+        Self {
+            db,
+            cached_vendored_root,
+        }
     }
 
     /// Map a `ResolvedDefinition` from a stub file to corresponding definitions in source files.
@@ -24,7 +32,9 @@ impl<'db> StubMapper<'db> {
         &self,
         def: ResolvedDefinition<'db>,
     ) -> impl Iterator<Item = ResolvedDefinition<'db>> {
-        if let Some(definitions) = map_stub_definition(self.db, &def) {
+        if let Some(definitions) =
+            map_stub_definition(self.db, &def, self.cached_vendored_root.as_deref())
+        {
             return Either::Left(definitions.into_iter());
         }
         Either::Right(std::iter::once(def))

crates/ty_python_semantic/src/types/ide_support.rs

@@ -1138,8 +1138,10 @@ mod resolve_definition {
     }
 
     use indexmap::IndexSet;
-    use ruff_db::files::{File, FileRange};
+    use ruff_db::files::{File, FileRange, vendored_path_to_file};
     use ruff_db::parsed::{ParsedModuleRef, parsed_module};
+    use ruff_db::system::SystemPath;
+    use ruff_db::vendored::VendoredPathBuf;
     use ruff_python_ast as ast;
     use rustc_hash::FxHashSet;
     use tracing::trace;
@@ -1397,17 +1399,42 @@ mod resolve_definition {
     pub fn map_stub_definition<'db>(
         db: &'db dyn Db,
         def: &ResolvedDefinition<'db>,
+        cached_vendored_typeshed: Option<&SystemPath>,
     ) -> 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);
+        trace!("Stub mapping definition in: {}", stub_file.path(db));
         if !stub_file.is_stub(db) {
             trace!("File isn't a stub, no stub mapping to do");
             return None;
         }
 
+        // We write vendored typeshed stubs to disk in the cache, and consequently "forget"
+        // that they're typeshed when an IDE hands those paths back to us later. For most
+        // purposes this seemingly doesn't matter at all, and avoids issues with someone
+        // editing the cache by hand in their IDE and us getting confused about the contents
+        // of the file (hello and welcome to anyone who has found Bigger Issues this causes).
+        //
+        // The major exception is in exactly stub-mapping, where we need to "remember" that
+        // we're in typeshed to successfully stub-map to the Real Stdlib. So here we attempt
+        // to do just that. The resulting file must not be used for anything other than
+        // this module lookup, as the `ResolvedDefinition` we're handling isn't for that file.
+        let mut stub_file_for_module_lookup = stub_file;
+        if let Some(vendored_typeshed) = cached_vendored_typeshed
+            && let Some(stub_path) = stub_file.path(db).as_system_path()
+            && let Ok(rel_path) = stub_path.strip_prefix(vendored_typeshed)
+            && let Ok(typeshed_file) =
+                vendored_path_to_file(db, VendoredPathBuf::from(rel_path.as_str()))
+        {
+            trace!(
+                "Stub is cached vendored typeshed: {}",
+                typeshed_file.path(db)
+            );
+            stub_file_for_module_lookup = typeshed_file;
+        }
+
         // It's definitely a stub, so now rerun module resolution but with stubs disabled.
-        let stub_module = file_to_module(db, stub_file)?;
+        let stub_module = file_to_module(db, stub_file_for_module_lookup)?;
         trace!("Found stub module: {}", stub_module.name(db));
         let real_module = resolve_real_module(db, stub_module.name(db))?;
         trace!("Found real module: {}", real_module.name(db));

crates/ty_server/Cargo.toml

@@ -22,7 +22,6 @@ ty_combine = { workspace = true }
 ty_ide = { workspace = true }
 ty_project = { workspace = true }
 ty_python_semantic = { workspace = true }
-ty_vendored = { workspace = true }
 
 anyhow = { workspace = true }
 bitflags = { workspace = true }

crates/ty_server/src/system.rs

@@ -13,6 +13,7 @@ use ruff_db::system::{
     SystemPath, SystemPathBuf, SystemVirtualPath, SystemVirtualPathBuf, WritableSystem,
 };
 use ruff_notebook::{Notebook, NotebookError};
+use ty_ide::cached_vendored_path;
 use ty_python_semantic::Db;
 
 use crate::DocumentQuery;
@@ -25,20 +26,7 @@ pub(crate) fn file_to_url(db: &dyn Db, file: File) -> Option<Url> {
         FilePath::System(system) => Url::from_file_path(system.as_std_path()).ok(),
         FilePath::SystemVirtual(path) => Url::parse(path.as_str()).ok(),
         FilePath::Vendored(path) => {
-            let writable = db.system().as_writable()?;
-
-            let system_path = SystemPathBuf::from(format!(
-                "vendored/typeshed/{}/{}",
-                // The vendored files are uniquely identified by the source commit.
-                ty_vendored::SOURCE_COMMIT,
-                path.as_str()
-            ));
-
-            // Extract the vendored file onto the system.
-            let system_path = writable
-                .get_or_cache(&system_path, &|| db.vendored().read_to_string(path))
-                .ok()
-                .flatten()?;
+            let system_path = cached_vendored_path(db, path)?;
 
             Url::from_file_path(system_path.as_std_path()).ok()
         }