[red-knot] Move module-resolution logic to its own crate (#11964)

Cargo.lock

@@ -1978,7 +1978,7 @@ dependencies = [
  "notify",
  "parking_lot",
  "rayon",
- "red_knot_python_semantic",
+ "red_knot_module_resolver",
  "ruff_index",
  "ruff_notebook",
  "ruff_python_ast",
@@ -1995,6 +1995,19 @@ dependencies = [
  "zip",
 ]
 
+[[package]]
+name = "red_knot_module_resolver"
+version = "0.0.0"
+dependencies = [
+ "anyhow",
+ "ruff_db",
+ "ruff_python_stdlib",
+ "salsa",
+ "smol_str",
+ "tempfile",
+ "tracing",
+]
+
 [[package]]
 name = "red_knot_python_semantic"
 version = "0.0.0"
@@ -2003,17 +2016,16 @@ dependencies = [
  "bitflags 2.5.0",
  "hashbrown 0.14.5",
  "indexmap",
+ "red_knot_module_resolver",
  "ruff_db",
  "ruff_index",
  "ruff_python_ast",
  "ruff_python_parser",
- "ruff_python_stdlib",
  "ruff_text_size",
  "rustc-hash",
  "salsa",
  "smallvec",
  "smol_str",
- "tempfile",
  "tracing",
 ]
 

Cargo.toml

@@ -35,7 +35,7 @@ ruff_source_file = { path = "crates/ruff_source_file" }
 ruff_text_size = { path = "crates/ruff_text_size" }
 ruff_workspace = { path = "crates/ruff_workspace" }
 
-red_knot_python_semantic = { path = "crates/red_knot_python_semantic" }
+red_knot_module_resolver = { path = "crates/red_knot_module_resolver" }
 
 aho-corasick = { version = "1.1.3" }
 annotate-snippets = { version = "0.9.2", features = ["color"] }

crates/red_knot/Cargo.toml

@@ -12,7 +12,7 @@ license.workspace = true
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-red_knot_python_semantic = { workspace = true }
+red_knot_module_resolver = { workspace = true }
 
 ruff_python_parser = { workspace = true }
 ruff_python_ast = { workspace = true }

crates/red_knot/src/module.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 use dashmap::mapref::entry::Entry;
 use smol_str::SmolStr;
 
-use red_knot_python_semantic::module::ModuleKind;
+use red_knot_module_resolver::ModuleKind;
 
 use crate::db::{QueryResult, SemanticDb, SemanticJar};
 use crate::files::FileId;

crates/red_knot_module_resolver/Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+name = "red_knot_module_resolver"
+version = "0.0.0"
+publish = false
+authors = { workspace = true }
+edition = { workspace = true }
+rust-version = { workspace = true }
+homepage = { workspace = true }
+documentation = { workspace = true }
+repository = { workspace = true }
+license = { workspace = true }
+
+[dependencies]
+ruff_db = { workspace = true }
+ruff_python_stdlib = { workspace = true }
+
+salsa = { workspace = true }
+smol_str = { workspace = true }
+tracing = { workspace = true }
+
+[dev-dependencies]
+anyhow = { workspace = true }
+tempfile = { workspace = true }
+
+[lints]
+workspace = true

crates/red_knot_module_resolver/src/db.rs

@@ -0,0 +1,156 @@
+use ruff_db::Upcast;
+
+use crate::resolver::{
+    file_to_module,
+    internal::{ModuleNameIngredient, ModuleResolverSearchPaths},
+    resolve_module_query,
+};
+
+#[salsa::jar(db=Db)]
+pub struct Jar(
+    ModuleNameIngredient<'_>,
+    ModuleResolverSearchPaths,
+    resolve_module_query,
+    file_to_module,
+);
+
+pub trait Db: salsa::DbWithJar<Jar> + ruff_db::Db + Upcast<dyn ruff_db::Db> {}
+
+pub(crate) mod tests {
+    use std::sync;
+
+    use salsa::DebugWithDb;
+
+    use ruff_db::file_system::{FileSystem, MemoryFileSystem, OsFileSystem};
+    use ruff_db::vfs::Vfs;
+
+    use super::*;
+
+    #[salsa::db(Jar, ruff_db::Jar)]
+    pub(crate) struct TestDb {
+        storage: salsa::Storage<Self>,
+        file_system: TestFileSystem,
+        events: sync::Arc<sync::Mutex<Vec<salsa::Event>>>,
+        vfs: Vfs,
+    }
+
+    impl TestDb {
+        #[allow(unused)]
+        pub(crate) fn new() -> Self {
+            Self {
+                storage: salsa::Storage::default(),
+                file_system: TestFileSystem::Memory(MemoryFileSystem::default()),
+                events: sync::Arc::default(),
+                vfs: Vfs::with_stubbed_vendored(),
+            }
+        }
+
+        /// Returns the memory file system.
+        ///
+        /// ## Panics
+        /// If this test db isn't using a memory file system.
+        #[allow(unused)]
+        pub(crate) fn memory_file_system(&self) -> &MemoryFileSystem {
+            if let TestFileSystem::Memory(fs) = &self.file_system {
+                fs
+            } else {
+                panic!("The test db is not using a memory file system");
+            }
+        }
+
+        /// Uses the real file system instead of the memory file system.
+        ///
+        /// This useful for testing advanced file system features like permissions, symlinks, etc.
+        ///
+        /// Note that any files written to the memory file system won't be copied over.
+        #[allow(unused)]
+        pub(crate) fn with_os_file_system(&mut self) {
+            self.file_system = TestFileSystem::Os(OsFileSystem);
+        }
+
+        #[allow(unused)]
+        pub(crate) fn vfs_mut(&mut self) -> &mut Vfs {
+            &mut self.vfs
+        }
+
+        /// Takes the salsa events.
+        ///
+        /// ## Panics
+        /// If there are any pending salsa snapshots.
+        #[allow(unused)]
+        pub(crate) fn take_salsa_events(&mut self) -> Vec<salsa::Event> {
+            let inner = sync::Arc::get_mut(&mut self.events).expect("no pending salsa snapshots");
+
+            let events = inner.get_mut().unwrap();
+            std::mem::take(&mut *events)
+        }
+
+        /// Clears the salsa events.
+        ///
+        /// ## Panics
+        /// If there are any pending salsa snapshots.
+        #[allow(unused)]
+        pub(crate) fn clear_salsa_events(&mut self) {
+            self.take_salsa_events();
+        }
+    }
+
+    impl Upcast<dyn ruff_db::Db> for TestDb {
+        fn upcast(&self) -> &(dyn ruff_db::Db + 'static) {
+            self
+        }
+    }
+
+    impl ruff_db::Db for TestDb {
+        fn file_system(&self) -> &dyn ruff_db::file_system::FileSystem {
+            self.file_system.inner()
+        }
+
+        fn vfs(&self) -> &ruff_db::vfs::Vfs {
+            &self.vfs
+        }
+    }
+
+    impl Db for TestDb {}
+
+    impl salsa::Database for TestDb {
+        fn salsa_event(&self, event: salsa::Event) {
+            tracing::trace!("event: {:?}", event.debug(self));
+            let mut events = self.events.lock().unwrap();
+            events.push(event);
+        }
+    }
+
+    impl salsa::ParallelDatabase for TestDb {
+        fn snapshot(&self) -> salsa::Snapshot<Self> {
+            salsa::Snapshot::new(Self {
+                storage: self.storage.snapshot(),
+                file_system: self.file_system.snapshot(),
+                events: self.events.clone(),
+                vfs: self.vfs.snapshot(),
+            })
+        }
+    }
+
+    enum TestFileSystem {
+        Memory(MemoryFileSystem),
+        #[allow(unused)]
+        Os(OsFileSystem),
+    }
+
+    impl TestFileSystem {
+        fn inner(&self) -> &dyn FileSystem {
+            match self {
+                Self::Memory(inner) => inner,
+                Self::Os(inner) => inner,
+            }
+        }
+
+        fn snapshot(&self) -> Self {
+            match self {
+                Self::Memory(inner) => Self::Memory(inner.snapshot()),
+                Self::Os(inner) => Self::Os(inner.snapshot()),
+            }
+        }
+    }
+}

crates/red_knot_module_resolver/src/lib.rs

@@ -0,0 +1,7 @@
+mod db;
+mod module;
+mod resolver;
+
+pub use db::{Db, Jar};
+pub use module::{ModuleKind, ModuleName};
+pub use resolver::{resolve_module, set_module_resolution_settings, ModuleResolutionSettings};

crates/red_knot_module_resolver/src/module.rs

@@ -8,8 +8,6 @@ use ruff_python_stdlib::identifiers::is_identifier;
 
 use crate::Db;
 
-pub mod resolver;
-
 /// A module name, e.g. `foo.bar`.
 ///
 /// Always normalized to the absolute form (never a relative module name, i.e., never `.foo`).
@@ -46,7 +44,7 @@ impl ModuleName {
     /// ## Examples
     ///
     /// ```
-    /// use red_knot_python_semantic::module::ModuleName;
+    /// use red_knot_module_resolver::ModuleName;
     ///
     /// assert_eq!(ModuleName::new_static("foo.bar").as_deref(), Some("foo.bar"));
     /// assert_eq!(ModuleName::new_static(""), None);
@@ -78,7 +76,7 @@ impl ModuleName {
     /// # Examples
     ///
     /// ```
-    /// use red_knot_python_semantic::module::ModuleName;
+    /// use red_knot_module_resolver::ModuleName;
     ///
     /// assert_eq!(ModuleName::new_static("foo.bar.baz").unwrap().components().collect::<Vec<_>>(), vec!["foo", "bar", "baz"]);
     /// ```
@@ -91,7 +89,7 @@ impl ModuleName {
     /// # Examples
     ///
     /// ```
-    /// use red_knot_python_semantic::module::ModuleName;
+    /// use red_knot_module_resolver::ModuleName;
     ///
     /// assert_eq!(ModuleName::new_static("foo.bar").unwrap().parent(), Some(ModuleName::new_static("foo").unwrap()));
     /// assert_eq!(ModuleName::new_static("foo.bar.baz").unwrap().parent(), Some(ModuleName::new_static("foo.bar").unwrap()));
@@ -110,7 +108,7 @@ impl ModuleName {
     /// # Examples
     ///
     /// ```
-    /// use red_knot_python_semantic::module::ModuleName;
+    /// use red_knot_module_resolver::ModuleName;
     ///
     /// assert!(ModuleName::new_static("foo.bar").unwrap().starts_with(&ModuleName::new_static("foo").unwrap()));
     ///
@@ -135,7 +133,7 @@ impl ModuleName {
         &self.0
     }
 
-    fn from_relative_path(path: &FileSystemPath) -> Option<Self> {
+    pub(crate) fn from_relative_path(path: &FileSystemPath) -> Option<Self> {
         let path = if path.ends_with("__init__.py") || path.ends_with("__init__.pyi") {
             path.parent()?
         } else {
@@ -196,6 +194,22 @@ pub struct Module {
 }
 
 impl Module {
+    pub(crate) fn new(
+        name: ModuleName,
+        kind: ModuleKind,
+        search_path: ModuleSearchPath,
+        file: VfsFile,
+    ) -> Self {
+        Self {
+            inner: Arc::new(ModuleInner {
+                name,
+                kind,
+                search_path,
+                file,
+            }),
+        }
+    }
+
     /// The absolute name of the module (e.g. `foo.bar`)
     pub fn name(&self) -> &ModuleName {
         &self.inner.name

crates/red_knot_module_resolver/src/resolver.rs

@@ -1,14 +1,11 @@
 use salsa::DebugWithDb;
 use std::ops::Deref;
-use std::sync::Arc;
 
 use ruff_db::file_system::{FileSystem, FileSystemPath, FileSystemPathBuf};
 use ruff_db::vfs::{system_path_to_file, vfs_path_to_file, VfsFile, VfsPath};
 
-use crate::module::resolver::internal::ModuleResolverSearchPaths;
+use crate::module::{Module, ModuleKind, ModuleName, ModuleSearchPath, ModuleSearchPathKind};
-use crate::module::{
+use crate::resolver::internal::ModuleResolverSearchPaths;
-    Module, ModuleInner, ModuleKind, ModuleName, ModuleSearchPath, ModuleSearchPathKind,
-};
 use crate::Db;
 
 const TYPESHED_STDLIB_DIRECTORY: &str = "stdlib";
@@ -51,14 +48,7 @@ pub(crate) fn resolve_module_query<'db>(
 
     let (search_path, module_file, kind) = resolve_name(db, name)?;
 
-    let module = Module {
+    let module = Module::new(name.clone(), kind, search_path, module_file);
-        inner: Arc::new(ModuleInner {
-            name: name.clone(),
-            kind,
-            search_path,
-            file: module_file,
-        }),
-    };
 
     Some(module)
 }
@@ -84,6 +74,7 @@ pub fn path_to_module(db: &dyn Db, path: &VfsPath) -> Option<Module> {
 ///
 /// Returns `None` if the file is not a module locatable via `sys.path`.
 #[salsa::tracked]
+#[allow(unused)]
 pub(crate) fn file_to_module(db: &dyn Db, file: VfsFile) -> Option<Module> {
     let _ = tracing::trace_span!("file_to_module", file = ?file.debug(db.upcast())).enter();
 
@@ -127,7 +118,7 @@ pub(crate) fn file_to_module(db: &dyn Db, file: VfsFile) -> Option<Module> {
     }
 }
 
-/// Configures the [`ModuleSearchPath`]s that are used to resolve modules.
+/// Configures the search paths that are used to resolve modules.
 #[derive(Eq, PartialEq, Debug)]
 pub struct ModuleResolutionSettings {
     /// List of user-provided paths that should take first priority in the module resolution.
@@ -208,8 +199,8 @@ impl Deref for OrderedSearchPaths {
 // TODO(micha): Contribute a fix for this upstream where the singleton methods have the same visibility as the struct.
 #[allow(unreachable_pub, clippy::used_underscore_binding)]
 pub(crate) mod internal {
-    use crate::module::resolver::OrderedSearchPaths;
     use crate::module::ModuleName;
+    use crate::resolver::OrderedSearchPaths;
 
     #[salsa::input(singleton)]
     pub(crate) struct ModuleResolverSearchPaths {

crates/red_knot_python_semantic/Cargo.toml

@@ -11,10 +11,10 @@ repository = { workspace = true }
 license = { workspace = true }
 
 [dependencies]
+red_knot_module_resolver = { workspace = true }
 ruff_db = { workspace = true }
 ruff_index = { workspace = true }
 ruff_python_ast = { workspace = true }
-ruff_python_stdlib = { workspace = true }
 ruff_text_size = { workspace = true }
 
 bitflags = { workspace = true }
@@ -29,7 +29,6 @@ hashbrown = { workspace = true }
 [dev-dependencies]
 anyhow = { workspace = true }
 ruff_python_parser = { workspace = true }
-tempfile = { workspace = true }
 
 [lints]
 workspace = true

crates/red_knot_python_semantic/src/db.rs

@@ -2,10 +2,7 @@ use salsa::DbWithJar;
 
 use ruff_db::{Db as SourceDb, Upcast};
 
-use crate::module::resolver::{
+use red_knot_module_resolver::Db as ResolverDb;
-    file_to_module, internal::ModuleNameIngredient, internal::ModuleResolverSearchPaths,
-    resolve_module_query,
-};
 
 use crate::semantic_index::symbol::{public_symbols_map, scopes_map, PublicSymbolId, ScopeId};
 use crate::semantic_index::{root_scope, semantic_index, symbol_table};
@@ -13,13 +10,9 @@ use crate::types::{infer_types, public_symbol_ty};
 
 #[salsa::jar(db=Db)]
 pub struct Jar(
-    ModuleNameIngredient<'_>,
-    ModuleResolverSearchPaths,
     ScopeId<'_>,
     PublicSymbolId<'_>,
     symbol_table,
-    resolve_module_query,
-    file_to_module,
     scopes_map,
     root_scope,
     semantic_index,
@@ -29,7 +22,10 @@ pub struct Jar(
 );
 
 /// Database giving access to semantic information about a Python program.
-pub trait Db: SourceDb + DbWithJar<Jar> + Upcast<dyn SourceDb> {}
+pub trait Db:
+    SourceDb + ResolverDb + DbWithJar<Jar> + Upcast<dyn SourceDb> + Upcast<dyn ResolverDb>
+{
+}
 
 #[cfg(test)]
 pub(crate) mod tests {
@@ -42,13 +38,14 @@ pub(crate) mod tests {
     use salsa::storage::HasIngredientsFor;
     use salsa::DebugWithDb;
 
+    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};
 
     use super::{Db, Jar};
 
-    #[salsa::db(Jar, SourceJar)]
+    #[salsa::db(Jar, ResolverJar, SourceJar)]
     pub(crate) struct TestDb {
         storage: salsa::Storage<Self>,
         vfs: Vfs,
@@ -78,15 +75,6 @@ pub(crate) mod tests {
             }
         }
 
-        /// Uses the real file system instead of the memory file system.
-        ///
-        /// This useful for testing advanced file system features like permissions, symlinks, etc.
-        ///
-        /// Note that any files written to the memory file system won't be copied over.
-        pub(crate) fn with_os_file_system(&mut self) {
-            self.file_system = TestFileSystem::Os(OsFileSystem);
-        }
-
         #[allow(unused)]
         pub(crate) fn vfs_mut(&mut self) -> &mut Vfs {
             &mut self.vfs
@@ -131,6 +119,13 @@ pub(crate) mod tests {
         }
     }
 
+    impl Upcast<dyn ResolverDb> for TestDb {
+        fn upcast(&self) -> &(dyn ResolverDb + 'static) {
+            self
+        }
+    }
+
+    impl red_knot_module_resolver::Db for TestDb {}
     impl Db for TestDb {}
 
     impl salsa::Database for TestDb {
@@ -157,6 +152,7 @@ pub(crate) mod tests {
 
     enum TestFileSystem {
         Memory(MemoryFileSystem),
+        #[allow(dead_code)]
         Os(OsFileSystem),
     }
 

crates/red_knot_python_semantic/src/lib.rs

@@ -1,6 +1,5 @@
 pub mod ast_node_ref;
 mod db;
-pub mod module;
 pub mod name;
 mod node_key;
 pub mod semantic_index;

crates/red_knot_python_semantic/src/types.rs

@@ -513,9 +513,9 @@ mod tests {
     use crate::db::tests::{
         assert_will_not_run_function_query, assert_will_run_function_query, TestDb,
     };
-    use crate::module::resolver::{set_module_resolution_settings, ModuleResolutionSettings};
     use crate::semantic_index::root_scope;
     use crate::types::{expression_ty, infer_types, public_symbol_ty_by_name, TypingContext};
+    use red_knot_module_resolver::{set_module_resolution_settings, ModuleResolutionSettings};
 
     fn setup_db() -> TestDb {
         let mut db = TestDb::new();

crates/red_knot_python_semantic/src/types/infer.rs

@@ -2,13 +2,13 @@ use std::sync::Arc;
 
 use rustc_hash::FxHashMap;
 
+use red_knot_module_resolver::resolve_module;
+use red_knot_module_resolver::ModuleName;
 use ruff_db::vfs::VfsFile;
 use ruff_index::IndexVec;
 use ruff_python_ast as ast;
 use ruff_python_ast::{ExprContext, TypeParams};
 
-use crate::module::resolver::resolve_module;
-use crate::module::ModuleName;
 use crate::name::Name;
 use crate::semantic_index::ast_ids::{ScopeAstIdNode, ScopeExpressionId};
 use crate::semantic_index::definition::{Definition, ImportDefinition, ImportFromDefinition};
@@ -358,7 +358,7 @@ impl<'db> TypeInferenceBuilder<'db> {
             } = alias;
 
             let module_name = ModuleName::new(&name.id);
-            let module = module_name.and_then(|name| resolve_module(self.db, name));
+            let module = module_name.and_then(|name| resolve_module(self.db.upcast(), name));
             let module_ty = module
                 .map(|module| self.typing_context().module_ty(module.file()))
                 .unwrap_or(Type::Unknown);
@@ -384,7 +384,8 @@ impl<'db> TypeInferenceBuilder<'db> {
         let import_id = import.scope_ast_id(self.db, self.file_id, self.file_scope_id);
         let module_name = ModuleName::new(module.as_deref().expect("Support relative imports"));
 
-        let module = module_name.and_then(|module_name| resolve_module(self.db, module_name));
+        let module =
+            module_name.and_then(|module_name| resolve_module(self.db.upcast(), module_name));
         let module_ty = module
             .map(|module| self.typing_context().module_ty(module.file()))
             .unwrap_or(Type::Unknown);
@@ -694,9 +695,9 @@ mod tests {
     use ruff_db::vfs::system_path_to_file;
 
     use crate::db::tests::TestDb;
-    use crate::module::resolver::{set_module_resolution_settings, ModuleResolutionSettings};
     use crate::name::Name;
     use crate::types::{public_symbol_ty_by_name, Type, TypingContext};
+    use red_knot_module_resolver::{set_module_resolution_settings, ModuleResolutionSettings};
 
     fn setup_db() -> TestDb {
         let mut db = TestDb::new();