Use Salsa queries to load TypeshedVersions

crates/red_knot_module_resolver/src/db.rs

@@ -6,7 +6,7 @@ use crate::resolver::{
     resolve_module_query,
 };
 use crate::supported_py_version::TargetPyVersion;
-use crate::typeshed::TypeshedVersions;
+use crate::typeshed::parse_typeshed_versions;
 
 #[salsa::jar(db=Db)]
 pub struct Jar(
@@ -15,14 +15,12 @@ pub struct Jar(
     TargetPyVersion,
     resolve_module_query,
     file_to_module,
+    parse_typeshed_versions,
 );
 
-pub trait Db: salsa::DbWithJar<Jar> + ruff_db::Db + Upcast<dyn ruff_db::Db> {
-    fn typeshed_versions(&self) -> &TypeshedVersions;
-}
+pub trait Db: salsa::DbWithJar<Jar> + ruff_db::Db + Upcast<dyn ruff_db::Db> {}
 
 pub(crate) mod tests {
-    use std::str::FromStr;
     use std::sync;
 
     use salsa::DebugWithDb;
@@ -38,7 +36,6 @@ pub(crate) mod tests {
         file_system: TestFileSystem,
         events: sync::Arc<sync::Mutex<Vec<salsa::Event>>>,
         vfs: Vfs,
-        typeshed_versions: TypeshedVersions,
     }
 
     impl TestDb {
@@ -49,7 +46,6 @@ pub(crate) mod tests {
                 file_system: TestFileSystem::Memory(MemoryFileSystem::default()),
                 events: sync::Arc::default(),
                 vfs: Vfs::with_stubbed_vendored(),
-                typeshed_versions: TypeshedVersions::from_str("").unwrap(),
             }
         }
 
@@ -119,11 +115,7 @@ pub(crate) mod tests {
         }
     }
 
-    impl Db for TestDb {
-        fn typeshed_versions(&self) -> &TypeshedVersions {
-            &self.typeshed_versions
-        }
-    }
+    impl Db for TestDb {}
 
     impl salsa::Database for TestDb {
         fn salsa_event(&self, event: salsa::Event) {
@@ -140,7 +132,6 @@ pub(crate) mod tests {
                 file_system: self.file_system.snapshot(),
                 events: self.events.clone(),
                 vfs: self.vfs.snapshot(),
-                typeshed_versions: self.typeshed_versions.clone(),
             })
         }
     }

crates/red_knot_module_resolver/src/path.rs

@@ -1,12 +1,12 @@
 use std::iter::FusedIterator;
 
 use ruff_db::file_system::{FileSystemPath, FileSystemPathBuf};
-use ruff_db::vfs::VfsPath;
+use ruff_db::vfs::{system_path_to_file, VfsPath};
 
 use crate::module_name::ModuleName;
 use crate::supported_py_version::get_target_py_version;
-use crate::typeshed::TypeshedVersionsQueryResult;
-use crate::Db;
+use crate::typeshed::{parse_typeshed_versions, TypeshedVersionsQueryResult};
+use crate::{Db, TypeshedVersions};
 
 #[repr(transparent)]
 #[derive(Debug, PartialEq, Eq, Hash)]
@@ -139,13 +139,13 @@ impl ModuleResolutionPath {
     }
 
     #[must_use]
-    pub(crate) fn is_regular_package(&self, db: &dyn Db) -> bool {
-        ModuleResolutionPathRef::from(self).is_regular_package(db)
+    pub(crate) fn is_regular_package(&self, db: &dyn Db, search_path: &Self) -> bool {
+        ModuleResolutionPathRef::from(self).is_regular_package(db, search_path)
     }
 
     #[must_use]
-    pub(crate) fn is_directory(&self, db: &dyn Db) -> bool {
-        ModuleResolutionPathRef::from(self).is_directory(db)
+    pub(crate) fn is_directory(&self, db: &dyn Db, search_path: &Self) -> bool {
+        ModuleResolutionPathRef::from(self).is_directory(db, search_path)
     }
 
     #[must_use]
@@ -345,57 +345,67 @@ impl<'a> ModuleResolutionPathRef<'a> {
     }
 
     #[must_use]
-    pub(crate) fn is_directory(&self, db: &dyn Db) -> bool {
-        match self {
-            Self::Extra(ExtraPath(path)) => db.file_system().is_directory(path),
-            Self::FirstParty(FirstPartyPath(path)) => db.file_system().is_directory(path),
-            Self::SitePackages(SitePackagesPath(path)) => db.file_system().is_directory(path),
-            Self::StandardLibrary(StandardLibraryPath(path)) => {
+    fn load_typeshed_versions(db: &dyn Db, stdlib_root: &StandardLibraryPath) -> TypeshedVersions {
+        let StandardLibraryPath(stdlib_fs_path) = stdlib_root;
+        let versions_path = stdlib_fs_path.join("VERSIONS");
+        let versions_file = system_path_to_file(db.upcast(), versions_path).unwrap();
+        parse_typeshed_versions(db, versions_file)
+    }
+
+    #[must_use]
+    pub(crate) fn is_directory(&self, db: &dyn Db, search_path: impl Into<Self>) -> bool {
+        match (self, search_path.into()) {
+            (Self::Extra(ExtraPath(path)), Self::Extra(_)) => db.file_system().is_directory(path),
+            (Self::FirstParty(FirstPartyPath(path)), Self::FirstParty(_)) => db.file_system().is_directory(path),
+            (Self::SitePackages(SitePackagesPath(path)), Self::SitePackages(_)) => db.file_system().is_directory(path),
+            (Self::StandardLibrary(StandardLibraryPath(path)), Self::StandardLibrary(stdlib_root)) => {
                 let Some(module_name) = self.as_module_name() else {
                     return false;
                 };
-                match db
-                    .typeshed_versions()
-                    .query_module(&module_name, get_target_py_version(db))
-                {
+                let typeshed_versions = Self::load_typeshed_versions(db, stdlib_root);
+                match typeshed_versions.query_module(&module_name, get_target_py_version(db)) {
                     TypeshedVersionsQueryResult::Exists
                     | TypeshedVersionsQueryResult::MaybeExists => {
                         db.file_system().is_directory(path)
                     }
                     TypeshedVersionsQueryResult::DoesNotExist => false,
                 }
             }
+            (path, root) => unreachable!(
+                "The search path should always be the same variant as `self` (got: {path:?}, {root:?})"
+            )
         }
     }
 
     #[must_use]
-    pub(crate) fn is_regular_package(&self, db: &dyn Db) -> bool {
-        match self {
-            Self::Extra(ExtraPath(fs_path))
-            | Self::FirstParty(FirstPartyPath(fs_path))
-            | Self::SitePackages(SitePackagesPath(fs_path)) => {
+    pub(crate) fn is_regular_package(&self, db: &dyn Db, search_path: impl Into<Self>) -> bool {
+        match (self, search_path.into()) {
+            (Self::Extra(ExtraPath(fs_path)), Self::Extra(_))
+            | (Self::FirstParty(FirstPartyPath(fs_path)), Self::FirstParty(_))
+            | (Self::SitePackages(SitePackagesPath(fs_path)), Self::SitePackages(_)) => {
                 let file_system = db.file_system();
                 file_system.exists(&fs_path.join("__init__.py"))
                     || file_system.exists(&fs_path.join("__init__.pyi"))
             }
             // Unlike the other variants:
             // (1) Account for VERSIONS
             // (2) Only test for `__init__.pyi`, not `__init__.py`
-            Self::StandardLibrary(StandardLibraryPath(fs_path)) => {
+            (Self::StandardLibrary(StandardLibraryPath(fs_path)), Self::StandardLibrary(stdlib_root)) => {
                 let Some(module_name) = self.as_module_name() else {
                     return false;
                 };
-                match db
-                    .typeshed_versions()
-                    .query_module(&module_name, get_target_py_version(db))
-                {
+                let typeshed_versions = Self::load_typeshed_versions(db, stdlib_root);
+                match typeshed_versions.query_module(&module_name, get_target_py_version(db)) {
                     TypeshedVersionsQueryResult::Exists
                     | TypeshedVersionsQueryResult::MaybeExists => {
                         db.file_system().exists(&fs_path.join("__init__.pyi"))
                     }
                     TypeshedVersionsQueryResult::DoesNotExist => false,
                 }
             }
+            (path, root) => unreachable!(
+                "The search path should always be the same variant as `self` (got: {path:?}, {root:?})"
+            )
         }
     }
 

crates/red_knot_module_resolver/src/resolver.rs

@@ -241,7 +241,7 @@ fn resolve_name(
                 package_path.push(module_name);
 
                 // Must be a `__init__.pyi` or `__init__.py` or it isn't a package.
-                let kind = if package_path.is_directory(db) {
+                let kind = if package_path.is_directory(db, search_path) {
                     package_path.push("__init__");
                     ModuleKind::Package
                 } else {
@@ -301,11 +301,11 @@ where
     for folder in components {
         package_path.push(folder);
 
-        let is_regular_package = package_path.is_regular_package(db);
+        let is_regular_package = package_path.is_regular_package(db, module_search_path);
 
         if is_regular_package {
             in_namespace_package = false;
-        } else if package_path.is_directory(db) {
+        } else if package_path.is_directory(db, module_search_path) {
             // A directory without an `__init__.py` is a namespace package, continue with the next folder.
             in_namespace_package = true;
         } else if in_namespace_package {

crates/red_knot_module_resolver/src/supported_py_version.rs

@@ -1,5 +1,7 @@
 #![allow(clippy::used_underscore_binding)] // necessary for Salsa inputs
 #![allow(unreachable_pub)]
+#![allow(clippy::needless_lifetimes)]
+#![allow(clippy::clone_on_copy)]
 
 use crate::Db;
 
@@ -18,7 +20,6 @@ pub enum SupportedPyVersion {
 
 #[salsa::input(singleton)]
 pub(crate) struct TargetPyVersion {
-    #[return_ref]
     pub(crate) target_py_version: SupportedPyVersion,
 }
 
@@ -31,5 +32,5 @@ pub(crate) fn set_target_py_version(db: &mut dyn Db, target_version: SupportedPy
 }
 
 pub(crate) fn get_target_py_version(db: &dyn Db) -> SupportedPyVersion {
-    *TargetPyVersion::get(db).target_py_version(db)
+    TargetPyVersion::get(db).target_py_version(db)
 }

crates/red_knot_module_resolver/src/typeshed.rs

@@ -1,7 +1,7 @@
 mod versions;
 
 pub use versions::TypeshedVersions;
-pub(crate) use versions::TypeshedVersionsQueryResult;
+pub(crate) use versions::{parse_typeshed_versions, TypeshedVersionsQueryResult};
 
 #[cfg(test)]
 mod tests {

crates/red_knot_module_resolver/src/typeshed/versions.rs

@@ -5,11 +5,19 @@ use std::ops::{RangeFrom, RangeInclusive};
 use std::str::FromStr;
 use std::sync::Arc;
 
+use ruff_db::{source::source_text, vfs::VfsFile};
 use rustc_hash::FxHashMap;
 
+use crate::db::Db;
 use crate::module_name::ModuleName;
 use crate::supported_py_version::SupportedPyVersion;
 
+#[salsa::tracked]
+pub(crate) fn parse_typeshed_versions(db: &dyn Db, versions_file: VfsFile) -> TypeshedVersions {
+    let file_content = source_text(db.upcast(), versions_file);
+    file_content.parse().unwrap()
+}
+
 #[derive(Debug, PartialEq, Eq)]
 pub struct TypeshedVersionsParseError {
     line_number: NonZeroU16,

crates/red_knot_python_semantic/src/db.rs

@@ -31,15 +31,14 @@ pub trait Db:
 pub(crate) mod tests {
     use std::fmt::Formatter;
     use std::marker::PhantomData;
-    use std::str::FromStr;
     use std::sync::Arc;
 
     use salsa::id::AsId;
     use salsa::ingredient::Ingredient;
     use salsa::storage::HasIngredientsFor;
     use salsa::DebugWithDb;
 
-    use red_knot_module_resolver::{Db as ResolverDb, Jar as ResolverJar, TypeshedVersions};
+    use red_knot_module_resolver::{Db as ResolverDb, Jar as ResolverJar};
     use ruff_db::file_system::{FileSystem, MemoryFileSystem, OsFileSystem};
     use ruff_db::vfs::Vfs;
     use ruff_db::{Db as SourceDb, Jar as SourceJar, Upcast};
@@ -52,7 +51,6 @@ pub(crate) mod tests {
         vfs: Vfs,
         file_system: TestFileSystem,
         events: std::sync::Arc<std::sync::Mutex<Vec<salsa::Event>>>,
-        typeshed_versions: TypeshedVersions,
     }
 
     impl TestDb {
@@ -62,7 +60,6 @@ pub(crate) mod tests {
                 file_system: TestFileSystem::Memory(MemoryFileSystem::default()),
                 events: std::sync::Arc::default(),
                 vfs: Vfs::with_stubbed_vendored(),
-                typeshed_versions: TypeshedVersions::from_str("").unwrap(),
             }
         }
 
@@ -128,11 +125,7 @@ pub(crate) mod tests {
         }
     }
 
-    impl red_knot_module_resolver::Db for TestDb {
-        fn typeshed_versions(&self) -> &TypeshedVersions {
-            &self.typeshed_versions
-        }
-    }
+    impl red_knot_module_resolver::Db for TestDb {}
     impl Db for TestDb {}
 
     impl salsa::Database for TestDb {
@@ -153,7 +146,6 @@ pub(crate) mod tests {
                     TestFileSystem::Os(fs) => TestFileSystem::Os(fs.snapshot()),
                 },
                 events: self.events.clone(),
-                typeshed_versions: self.typeshed_versions.clone(),
             })
         }
     }