Teach file_to_module that vendored paths are equivalent to their cached equivalent

crates/ty_python_semantic/Cargo.toml

@@ -24,6 +24,7 @@ ruff_text_size = { workspace = true }
 ruff_python_literal = { workspace = true }
 ruff_python_trivia = { workspace = true }
 ty_static = { workspace = true }
+ty_vendored = { workspace = true }
 
 anyhow = { workspace = true }
 bitflags = { workspace = true }
@@ -58,7 +59,6 @@ ruff_python_parser = { workspace = true }
 ty_python_semantic = { workspace = true, features = ["testing"] }
 ty_static = { workspace = true }
 ty_test = { workspace = true }
-ty_vendored = { workspace = true }
 
 anyhow = { workspace = true }
 dir-test = { workspace = true }

crates/ty_python_semantic/src/lib.rs

@@ -13,7 +13,7 @@ pub use diagnostic::add_inferred_python_version_hint_to_diagnostic;
 pub use module_name::{ModuleName, ModuleNameResolutionError};
 pub use module_resolver::{
     Module, SearchPath, SearchPathValidationError, SearchPaths, all_modules, list_modules,
-    resolve_module, resolve_real_module, system_module_search_paths,
+    resolve_module, resolve_real_module, system_module_search_paths, vendored_path_for_cache,
 };
 pub use program::{
     Program, ProgramSettings, PythonVersionFileSource, PythonVersionSource,

crates/ty_python_semantic/src/module_resolver/mod.rs

@@ -3,7 +3,7 @@ use std::iter::FusedIterator;
 pub use list::{all_modules, list_modules};
 pub(crate) use module::KnownModule;
 pub use module::Module;
-pub use path::{SearchPath, SearchPathValidationError};
+pub use path::{SearchPath, SearchPathValidationError, vendored_path_for_cache};
 pub use resolver::SearchPaths;
 pub(crate) use resolver::file_to_module;
 pub use resolver::{resolve_module, resolve_real_module};

crates/ty_python_semantic/src/module_resolver/path.rs

@@ -607,7 +607,11 @@ impl SearchPath {
     }
 
     #[must_use]
-    pub(crate) fn relativize_system_path(&self, path: &SystemPath) -> Option<ModulePath> {
+    pub(crate) fn relativize_system_path(
+        &self,
+        db: &dyn Db,
+        path: &SystemPath,
+    ) -> Option<ModulePath> {
         if path
             .extension()
             .is_some_and(|extension| !self.is_valid_extension(extension))
@@ -629,7 +633,15 @@ impl SearchPath {
                         relative_path: relative_path.as_utf8_path().to_path_buf(),
                     })
             }
-            SearchPathInner::StandardLibraryVendored(_) => None,
+            // Check if this system path is actually a cached vendored path
+            SearchPathInner::StandardLibraryVendored(search_path) => path
+                .as_utf8_path()
+                .strip_prefix(cached_vendored_path(db, search_path)?.as_utf8_path())
+                .ok()
+                .map(|relative_path| ModulePath {
+                    search_path: self.clone(),
+                    relative_path: relative_path.to_path_buf(),
+                }),
         }
     }
 
@@ -721,6 +733,38 @@ impl SearchPath {
     }
 }
 
+/// Get the cache-relative `SystemPath` that a `VendoredPath` refers to
+pub fn vendored_path_for_cache(path: &VendoredPath) -> SystemPathBuf {
+    SystemPathBuf::from(format!(
+        "vendored/typeshed/{}/{}",
+        // The vendored files are uniquely identified by the source commit.
+        ty_vendored::SOURCE_COMMIT,
+        path.as_str()
+    ))
+}
+
+/// Get the full `SystemPath` in the cache that a `VendoredPath` refers to
+pub(crate) fn cached_vendored_path(db: &dyn Db, path: &VendoredPath) -> Option<SystemPathBuf> {
+    let writable = db.system().as_writable()?;
+    let system_path = vendored_path_for_cache(path);
+    Some(writable.cache_dir()?.join(system_path))
+}
+
+/// `==` but with support for treating vendored paths as equivalent to their cached system paths.
+pub(crate) fn path_is_equivalent(db: &dyn Db, path1: &FilePath, path2: &FilePath) -> bool {
+    if path1 == path2 {
+        true
+    } else if let (FilePath::System(system), FilePath::Vendored(vendored))
+    | (FilePath::Vendored(vendored), FilePath::System(system)) = (path1, path2)
+    {
+        cached_vendored_path(db, vendored)
+            .map(|cached| &cached == system)
+            .unwrap_or(false)
+    } else {
+        false
+    }
+}
+
 impl PartialEq<SystemPath> for SearchPath {
     fn eq(&self, other: &SystemPath) -> bool {
         self.as_system_path().is_some_and(|path| path == other)
@@ -1060,13 +1104,19 @@ mod tests {
 
         // Must have a `.pyi` extension or no extension:
         let bad_absolute_path = SystemPath::new("foo/stdlib/x.py");
-        assert_eq!(root.relativize_system_path(bad_absolute_path), None);
+        assert_eq!(root.relativize_system_path(&db, bad_absolute_path), None);
         let second_bad_absolute_path = SystemPath::new("foo/stdlib/x.rs");
-        assert_eq!(root.relativize_system_path(second_bad_absolute_path), None);
+        assert_eq!(
+            root.relativize_system_path(&db, second_bad_absolute_path),
+            None
+        );
 
         // Must be a path that is a child of `root`:
         let third_bad_absolute_path = SystemPath::new("bar/stdlib/x.pyi");
-        assert_eq!(root.relativize_system_path(third_bad_absolute_path), None);
+        assert_eq!(
+            root.relativize_system_path(&db, third_bad_absolute_path),
+            None
+        );
     }
 
     #[test]
@@ -1076,10 +1126,13 @@ mod tests {
         let root = SearchPath::extra(db.system(), src.clone()).unwrap();
         // Must have a `.py` extension, a `.pyi` extension, or no extension:
         let bad_absolute_path = src.join("x.rs");
-        assert_eq!(root.relativize_system_path(&bad_absolute_path), None);
+        assert_eq!(root.relativize_system_path(&db, &bad_absolute_path), None);
         // Must be a path that is a child of `root`:
         let second_bad_absolute_path = SystemPath::new("bar/src/x.pyi");
-        assert_eq!(root.relativize_system_path(second_bad_absolute_path), None);
+        assert_eq!(
+            root.relativize_system_path(&db, second_bad_absolute_path),
+            None
+        );
     }
 
     #[test]
@@ -1088,7 +1141,7 @@ mod tests {
         let src_search_path = SearchPath::first_party(db.system(), src.clone()).unwrap();
         let eggs_package = src.join("eggs/__init__.pyi");
         let module_path = src_search_path
-            .relativize_system_path(&eggs_package)
+            .relativize_system_path(&db, &eggs_package)
             .unwrap();
         assert_eq!(
             &module_path.relative_path,

crates/ty_python_semantic/src/module_resolver/resolver.rs

@@ -26,6 +26,7 @@ use ruff_python_ast::{
 
 use crate::db::Db;
 use crate::module_name::ModuleName;
+use crate::module_resolver::path::path_is_equivalent;
 use crate::module_resolver::typeshed::{TypeshedVersions, vendored_typeshed_versions};
 use crate::{Program, SearchPathSettings};
 
@@ -141,10 +142,9 @@ pub(crate) fn file_to_module(db: &dyn Db, file: File) -> Option<Module<'_>> {
     let _span = tracing::trace_span!("file_to_module", ?file).entered();
 
     let path = SystemOrVendoredPathRef::try_from_file(db, file)?;
-
     let module_name = search_paths(db, ModuleResolveMode::StubsAllowed).find_map(|candidate| {
         let relative_path = match path {
-            SystemOrVendoredPathRef::System(path) => candidate.relativize_system_path(path),
+            SystemOrVendoredPathRef::System(path) => candidate.relativize_system_path(db, path),
             SystemOrVendoredPathRef::Vendored(path) => candidate.relativize_vendored_path(path),
         }?;
         relative_path.to_module_name()
@@ -157,7 +157,7 @@ pub(crate) fn file_to_module(db: &dyn Db, file: File) -> Option<Module<'_>> {
     let module = resolve_module(db, &module_name)?;
     let module_file = module.file(db)?;
 
-    if file.path(db) == module_file.path(db) {
+    if path_is_equivalent(db, file.path(db), module_file.path(db)) {
         return Some(module);
     } else if file.source_type(db) == PySourceType::Python
         && module_file.source_type(db) == PySourceType::Stub

crates/ty_server/src/system.rs

@@ -13,7 +13,7 @@ use ruff_db::system::{
     SystemPath, SystemPathBuf, SystemVirtualPath, SystemVirtualPathBuf, WritableSystem,
 };
 use ruff_notebook::{Notebook, NotebookError};
-use ty_python_semantic::Db;
+use ty_python_semantic::{Db, vendored_path_for_cache};
 
 use crate::DocumentQuery;
 use crate::document::DocumentKey;
@@ -26,21 +26,15 @@ pub(crate) fn file_to_url(db: &dyn Db, file: File) -> Option<Url> {
         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()
-            ));
+            let system_path = vendored_path_for_cache(path);
 
             // Extract the vendored file onto the system.
-            let system_path = writable
+            let path = writable
                 .get_or_cache(&system_path, &|| db.vendored().read_to_string(path))
                 .ok()
                 .flatten()?;
 
-            Url::from_file_path(system_path.as_std_path()).ok()
+            Url::from_file_path(path.as_std_path()).ok()
         }
     }
 }