[red-knot] Watch search paths (#12407)

Cargo.lock

@@ -1898,6 +1898,7 @@ version = "0.0.0"
 dependencies = [
  "anyhow",
  "bitflags 2.6.0",
+ "countme",
  "hashbrown",
  "ordermap",
  "red_knot_module_resolver",

crates/red_knot/src/main.rs

@@ -12,7 +12,7 @@ use tracing_tree::time::Uptime;
 
 use red_knot::db::RootDatabase;
 use red_knot::watch;
-use red_knot::watch::Watcher;
+use red_knot::watch::WorkspaceWatcher;
 use red_knot::workspace::WorkspaceMetadata;
 use ruff_db::program::{ProgramSettings, SearchPathSettings};
 use ruff_db::system::{OsSystem, System, SystemPathBuf};
@@ -142,7 +142,7 @@ struct MainLoop {
     receiver: crossbeam_channel::Receiver<MainLoopMessage>,
 
     /// The file system watcher, if running in watch mode.
-    watcher: Option<Watcher>,
+    watcher: Option<WorkspaceWatcher>,
 
     verbosity: Option<VerbosityLevel>,
 }
@@ -164,26 +164,23 @@ impl MainLoop {
 
     fn watch(mut self, db: &mut RootDatabase) -> anyhow::Result<()> {
         let sender = self.sender.clone();
-        let mut watcher = watch::directory_watcher(move |event| {
+        let watcher = watch::directory_watcher(move |event| {
             sender.send(MainLoopMessage::ApplyChanges(event)).unwrap();
         })?;
 
-        watcher.watch(db.workspace().root(db))?;
-
-        self.watcher = Some(watcher);
-
+        self.watcher = Some(WorkspaceWatcher::new(watcher, db));
         self.run(db);
 
         Ok(())
     }
 
     #[allow(clippy::print_stderr)]
-    fn run(self, db: &mut RootDatabase) {
+    fn run(mut self, db: &mut RootDatabase) {
         // Schedule the first check.
         self.sender.send(MainLoopMessage::CheckWorkspace).unwrap();
         let mut revision = 0usize;
 
-        for message in &self.receiver {
+        while let Ok(message) = self.receiver.recv() {
             tracing::trace!("Main Loop: Tick");
 
             match message {
@@ -224,6 +221,9 @@ impl MainLoop {
                     revision += 1;
                     // Automatically cancels any pending queries and waits for them to complete.
                     db.apply_changes(changes);
+                    if let Some(watcher) = self.watcher.as_mut() {
+                        watcher.update(db);
+                    }
                     self.sender.send(MainLoopMessage::CheckWorkspace).unwrap();
                 }
                 MainLoopMessage::Exit => {

crates/red_knot/src/watch.rs

@@ -1,7 +1,9 @@
 use ruff_db::system::{SystemPath, SystemPathBuf};
 pub use watcher::{directory_watcher, EventHandler, Watcher};
+pub use workspace_watcher::WorkspaceWatcher;
 
 mod watcher;
+mod workspace_watcher;
 
 /// Classification of a file system change event.
 ///

crates/red_knot/src/watch/workspace_watcher.rs

@@ -0,0 +1,112 @@
+use crate::db::RootDatabase;
+use crate::watch::Watcher;
+use ruff_db::system::SystemPathBuf;
+use rustc_hash::FxHashSet;
+use std::fmt::{Formatter, Write};
+use tracing::info;
+
+/// Wrapper around a [`Watcher`] that watches the relevant paths of a workspace.
+pub struct WorkspaceWatcher {
+    watcher: Watcher,
+
+    /// The paths that need to be watched. This includes paths for which setting up file watching failed.
+    watched_paths: FxHashSet<SystemPathBuf>,
+
+    /// Paths that should be watched but setting up the watcher failed for some reason.
+    /// This should be rare.
+    errored_paths: Vec<SystemPathBuf>,
+}
+
+impl WorkspaceWatcher {
+    /// Create a new workspace watcher.
+    pub fn new(watcher: Watcher, db: &RootDatabase) -> Self {
+        let mut watcher = Self {
+            watcher,
+            watched_paths: FxHashSet::default(),
+            errored_paths: Vec::new(),
+        };
+
+        watcher.update(db);
+
+        watcher
+    }
+
+    pub fn update(&mut self, db: &RootDatabase) {
+        let new_watch_paths = db.workspace().paths_to_watch(db);
+
+        let mut added_folders = new_watch_paths.difference(&self.watched_paths).peekable();
+        let mut removed_folders = self.watched_paths.difference(&new_watch_paths).peekable();
+
+        if added_folders.peek().is_none() && removed_folders.peek().is_none() {
+            return;
+        }
+
+        for added_folder in added_folders {
+            // Log a warning. It's not worth aborting if registering a single folder fails because
+            // Ruff otherwise stills works as expected.
+            if let Err(error) = self.watcher.watch(added_folder) {
+                // TODO: Log a user-facing warning.
+                tracing::warn!("Failed to setup watcher for path '{added_folder}': {error}. You have to restart Ruff after making changes to files under this path or you might see stale results.");
+                self.errored_paths.push(added_folder.clone());
+            }
+        }
+
+        for removed_path in removed_folders {
+            if let Some(index) = self
+                .errored_paths
+                .iter()
+                .position(|path| path == removed_path)
+            {
+                self.errored_paths.swap_remove(index);
+                continue;
+            }
+
+            if let Err(error) = self.watcher.unwatch(removed_path) {
+                info!("Failed to remove the file watcher for the path '{removed_path}: {error}.");
+            }
+        }
+
+        info!(
+            "Set up file watchers for {}",
+            DisplayWatchedPaths {
+                paths: &new_watch_paths
+            }
+        );
+
+        self.watched_paths = new_watch_paths;
+    }
+
+    /// Returns `true` if setting up watching for any path failed.
+    pub fn has_errored_paths(&self) -> bool {
+        !self.errored_paths.is_empty()
+    }
+
+    pub fn flush(&self) {
+        self.watcher.flush();
+    }
+
+    pub fn stop(self) {
+        self.watcher.stop();
+    }
+}
+
+struct DisplayWatchedPaths<'a> {
+    paths: &'a FxHashSet<SystemPathBuf>,
+}
+
+impl std::fmt::Display for DisplayWatchedPaths<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.write_char('[')?;
+
+        let mut iter = self.paths.iter();
+        if let Some(first) = iter.next() {
+            write!(f, "\"{first}\"")?;
+
+            for path in iter {
+                write!(f, ", \"{path}\"")?;
+            }
+        }
+
+        f.write_char(']')
+    }
+}

crates/red_knot/src/workspace.rs

@@ -6,6 +6,7 @@ use std::{collections::BTreeMap, sync::Arc};
 use rustc_hash::{FxBuildHasher, FxHashSet};
 
 pub use metadata::{PackageMetadata, WorkspaceMetadata};
+use red_knot_module_resolver::system_module_search_paths;
 use ruff_db::{
     files::{system_path_to_file, File},
     system::{walk_directory::WalkState, SystemPath, SystemPathBuf},
@@ -240,6 +241,17 @@ impl Workspace {
             FxHashSet::default()
         }
     }
+
+    /// Returns the paths that should be watched.
+    ///
+    /// The paths that require watching might change with every revision.
+    pub fn paths_to_watch(self, db: &dyn Db) -> FxHashSet<SystemPathBuf> {
+        ruff_db::system::deduplicate_nested_paths(
+            std::iter::once(self.root(db)).chain(system_module_search_paths(db.upcast())),
+        )
+        .map(SystemPath::to_path_buf)
+        .collect()
+    }
 }
 
 #[salsa::tracked]

crates/red_knot/tests/file_watching.rs

@@ -6,18 +6,18 @@ use anyhow::{anyhow, Context};
 
 use red_knot::db::RootDatabase;
 use red_knot::watch;
-use red_knot::watch::{directory_watcher, Watcher};
+use red_knot::watch::{directory_watcher, WorkspaceWatcher};
 use red_knot::workspace::WorkspaceMetadata;
 use red_knot_module_resolver::{resolve_module, ModuleName};
 use ruff_db::files::{system_path_to_file, File};
-use ruff_db::program::{ProgramSettings, SearchPathSettings, TargetVersion};
+use ruff_db::program::{Program, ProgramSettings, SearchPathSettings, TargetVersion};
 use ruff_db::source::source_text;
 use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf};
 use ruff_db::Upcast;
 
 struct TestCase {
     db: RootDatabase,
-    watcher: Option<Watcher>,
+    watcher: Option<WorkspaceWatcher>,
     changes_receiver: crossbeam::channel::Receiver<Vec<watch::ChangeEvent>>,
     temp_dir: tempfile::TempDir,
 }
@@ -55,6 +55,23 @@ impl TestCase {
         all_events
     }
 
+    fn update_search_path_settings(
+        &mut self,
+        f: impl FnOnce(&SearchPathSettings) -> SearchPathSettings,
+    ) {
+        let program = Program::get(self.db());
+        let search_path_settings = program.search_paths(self.db());
+
+        let new_settings = f(search_path_settings);
+
+        program.set_search_paths(&mut self.db).to(new_settings);
+
+        if let Some(watcher) = &mut self.watcher {
+            watcher.update(&self.db);
+            assert!(!watcher.has_errored_paths());
+        }
+    }
+
     fn collect_package_files(&self, path: &SystemPath) -> Vec<File> {
         let package = self.db().workspace().package(self.db(), path).unwrap();
         let files = package.files(self.db());
@@ -70,12 +87,37 @@ impl TestCase {
 }
 
 fn setup<I, P>(workspace_files: I) -> anyhow::Result<TestCase>
+where
+    I: IntoIterator<Item = (P, &'static str)>,
+    P: AsRef<SystemPath>,
+{
+    setup_with_search_paths(workspace_files, |_root, workspace_path| {
+        SearchPathSettings {
+            extra_paths: vec![],
+            workspace_root: workspace_path.to_path_buf(),
+            custom_typeshed: None,
+            site_packages: None,
+        }
+    })
+}
+
+fn setup_with_search_paths<I, P>(
+    workspace_files: I,
+    create_search_paths: impl FnOnce(&SystemPath, &SystemPath) -> SearchPathSettings,
+) -> anyhow::Result<TestCase>
 where
     I: IntoIterator<Item = (P, &'static str)>,
     P: AsRef<SystemPath>,
 {
     let temp_dir = tempfile::tempdir()?;
 
+    let root_path = SystemPath::from_std_path(temp_dir.path()).ok_or_else(|| {
+        anyhow!(
+            "Temp directory '{}' is not a valid UTF-8 path.",
+            temp_dir.path().display()
+        )
+    })?;
+
     let workspace_path = temp_dir.path().join("workspace");
 
     std::fs::create_dir_all(&workspace_path).with_context(|| {
@@ -96,7 +138,7 @@ where
         workspace_path
             .as_utf8_path()
             .canonicalize_utf8()
-            .with_context(|| "Failed to canonzialize workspace path.")?,
+            .with_context(|| "Failed to canonicalize workspace path.")?,
     );
 
     for (relative_path, content) in workspace_files {
@@ -115,25 +157,31 @@ where
     let system = OsSystem::new(&workspace_path);
 
     let workspace = WorkspaceMetadata::from_path(&workspace_path, &system)?;
+    let search_paths = create_search_paths(root_path, workspace.root());
+
+    for path in search_paths
+        .extra_paths
+        .iter()
+        .chain(search_paths.site_packages.iter())
+        .chain(search_paths.custom_typeshed.iter())
+    {
+        std::fs::create_dir_all(path.as_std_path())
+            .with_context(|| format!("Failed to create search path '{path}'"))?;
+    }
+
     let settings = ProgramSettings {
         target_version: TargetVersion::default(),
-        search_paths: SearchPathSettings {
-            extra_paths: vec![],
-            workspace_root: workspace.root().to_path_buf(),
-            custom_typeshed: None,
-            site_packages: None,
-        },
+        search_paths,
     };
 
     let db = RootDatabase::new(workspace, settings, system);
 
     let (sender, receiver) = crossbeam::channel::unbounded();
-    let mut watcher = directory_watcher(move |events| sender.send(events).unwrap())
+    let watcher = directory_watcher(move |events| sender.send(events).unwrap())
         .with_context(|| "Failed to create directory watcher")?;
 
-    watcher
-        .watch(&workspace_path)
-        .with_context(|| "Failed to set up watcher for workspace directory.")?;
+    let watcher = WorkspaceWatcher::new(watcher, &db);
+    assert!(!watcher.has_errored_paths());
 
     let test_case = TestCase {
         db,
@@ -597,3 +645,93 @@ fn directory_deleted() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn search_path() -> anyhow::Result<()> {
+    let mut case =
+        setup_with_search_paths([("bar.py", "import sub.a")], |root_path, workspace_path| {
+            SearchPathSettings {
+                extra_paths: vec![],
+                workspace_root: workspace_path.to_path_buf(),
+                custom_typeshed: None,
+                site_packages: Some(root_path.join("site_packages")),
+            }
+        })?;
+
+    let site_packages = case.root_path().join("site_packages");
+
+    assert_eq!(
+        resolve_module(case.db(), ModuleName::new("a").unwrap()),
+        None
+    );
+
+    std::fs::write(site_packages.join("a.py").as_std_path(), "class A: ...")?;
+    std::fs::write(site_packages.join("__init__.py").as_std_path(), "")?;
+
+    let changes = case.stop_watch();
+
+    case.db_mut().apply_changes(changes);
+
+    assert!(resolve_module(case.db().upcast(), ModuleName::new_static("a").unwrap()).is_some());
+    assert_eq!(
+        case.collect_package_files(&case.workspace_path("bar.py")),
+        &[case.system_file(case.workspace_path("bar.py")).unwrap()]
+    );
+
+    Ok(())
+}
+
+#[test]
+fn add_search_path() -> anyhow::Result<()> {
+    let mut case = setup([("bar.py", "import sub.a")])?;
+
+    let site_packages = case.workspace_path("site_packages");
+    std::fs::create_dir_all(site_packages.as_std_path())?;
+
+    assert!(resolve_module(case.db().upcast(), ModuleName::new_static("a").unwrap()).is_none());
+
+    // Register site-packages as a search path.
+    case.update_search_path_settings(|settings| SearchPathSettings {
+        site_packages: Some(site_packages.clone()),
+        ..settings.clone()
+    });
+
+    std::fs::write(site_packages.join("a.py").as_std_path(), "class A: ...")?;
+    std::fs::write(site_packages.join("__init__.py").as_std_path(), "")?;
+
+    let changes = case.stop_watch();
+
+    case.db_mut().apply_changes(changes);
+
+    assert!(resolve_module(case.db().upcast(), ModuleName::new_static("a").unwrap()).is_some());
+
+    Ok(())
+}
+
+#[test]
+fn remove_search_path() -> anyhow::Result<()> {
+    let mut case =
+        setup_with_search_paths([("bar.py", "import sub.a")], |root_path, workspace_path| {
+            SearchPathSettings {
+                extra_paths: vec![],
+                workspace_root: workspace_path.to_path_buf(),
+                custom_typeshed: None,
+                site_packages: Some(root_path.join("site_packages")),
+            }
+        })?;
+
+    // Remove site packages from the search path settings.
+    let site_packages = case.root_path().join("site_packages");
+    case.update_search_path_settings(|settings| SearchPathSettings {
+        site_packages: None,
+        ..settings.clone()
+    });
+
+    std::fs::write(site_packages.join("a.py").as_std_path(), "class A: ...")?;
+
+    let changes = case.stop_watch();
+
+    assert_eq!(changes, &[]);
+
+    Ok(())
+}

crates/red_knot_module_resolver/src/lib.rs

@@ -1,3 +1,16 @@
+use std::iter::FusedIterator;
+
+pub use db::{Db, Jar};
+pub use module::{Module, ModuleKind};
+pub use module_name::ModuleName;
+pub use resolver::resolve_module;
+use ruff_db::system::SystemPath;
+pub use typeshed::{
+    vendored_typeshed_stubs, TypeshedVersionsParseError, TypeshedVersionsParseErrorKind,
+};
+
+use crate::resolver::{module_resolution_settings, SearchPathIterator};
+
 mod db;
 mod module;
 mod module_name;
@@ -9,10 +22,29 @@ mod typeshed;
 #[cfg(test)]
 mod testing;
 
-pub use db::{Db, Jar};
-pub use module::{Module, ModuleKind};
-pub use module_name::ModuleName;
-pub use resolver::resolve_module;
-pub use typeshed::{
-    vendored_typeshed_stubs, TypeshedVersionsParseError, TypeshedVersionsParseErrorKind,
-};
+/// Returns an iterator over all search paths pointing to a system path
+pub fn system_module_search_paths(db: &dyn Db) -> SystemModuleSearchPathsIter {
+    SystemModuleSearchPathsIter {
+        inner: module_resolution_settings(db).search_paths(db),
+    }
+}
+
+pub struct SystemModuleSearchPathsIter<'db> {
+    inner: SearchPathIterator<'db>,
+}
+
+impl<'db> Iterator for SystemModuleSearchPathsIter<'db> {
+    type Item = &'db SystemPath;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            let next = self.inner.next()?;
+
+            if let Some(system_path) = next.as_system_path() {
+                return Some(system_path);
+            }
+        }
+    }
+}
+
+impl FusedIterator for SystemModuleSearchPathsIter<'_> {}

crates/red_knot_module_resolver/src/resolver.rs

@@ -258,7 +258,7 @@ pub(crate) fn editable_install_resolution_paths(db: &dyn Db) -> Vec<ModuleSearch
 /// are only calculated lazily.
 ///
 /// [`sys.path` at runtime]: https://docs.python.org/3/library/site.html#module-site
-struct SearchPathIterator<'db> {
+pub(crate) struct SearchPathIterator<'db> {
     db: &'db dyn Db,
     static_paths: std::slice::Iter<'db, ModuleSearchPath>,
     dynamic_paths: Option<std::slice::Iter<'db, ModuleSearchPath>>,
@@ -399,7 +399,7 @@ impl ModuleResolutionSettings {
         self.target_version
     }
 
-    fn search_paths<'db>(&'db self, db: &'db dyn Db) -> SearchPathIterator<'db> {
+    pub(crate) fn search_paths<'db>(&'db self, db: &'db dyn Db) -> SearchPathIterator<'db> {
         SearchPathIterator {
             db,
             static_paths: self.static_search_paths.iter(),

crates/red_knot_python_semantic/Cargo.toml

@@ -18,6 +18,7 @@ ruff_python_ast = { workspace = true }
 ruff_text_size = { workspace = true }
 
 bitflags = { workspace = true }
+countme = { workspace = true }
 ordermap = { workspace = true }
 salsa = { workspace = true }
 tracing = { workspace = true }

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -103,9 +103,13 @@ impl<'db> SemanticIndexBuilder<'db> {
 
         #[allow(unsafe_code)]
         // SAFETY: `node` is guaranteed to be a child of `self.module`
-        let scope_id = ScopeId::new(self.db, self.file, file_scope_id, unsafe {
-            node.to_kind(self.module.clone())
-        });
+        let scope_id = ScopeId::new(
+            self.db,
+            self.file,
+            file_scope_id,
+            unsafe { node.to_kind(self.module.clone()) },
+            countme::Count::default(),
+        );
 
         self.scope_ids_by_scope.push(scope_id);
         self.scopes_by_node.insert(node.node_key(), file_scope_id);
@@ -180,6 +184,7 @@ impl<'db> SemanticIndexBuilder<'db> {
             unsafe {
                 definition_node.into_owned(self.module.clone())
             },
+            countme::Count::default(),
         );
 
         self.definitions_by_node
@@ -201,6 +206,7 @@ impl<'db> SemanticIndexBuilder<'db> {
             unsafe {
                 AstNodeRef::new(self.module.clone(), expression_node)
             },
+            countme::Count::default(),
         );
         self.expressions_by_node
             .insert(expression_node.into(), expression);

crates/red_knot_python_semantic/src/semantic_index/definition.rs

@@ -24,6 +24,9 @@ pub struct Definition<'db> {
     #[no_eq]
     #[return_ref]
     pub(crate) node: DefinitionKind,
+
+    #[no_eq]
+    count: countme::Count<Definition<'static>>,
 }
 
 impl<'db> Definition<'db> {

crates/red_knot_python_semantic/src/semantic_index/expression.rs

@@ -22,6 +22,9 @@ pub(crate) struct Expression<'db> {
     #[no_eq]
     #[return_ref]
     pub(crate) node: AstNodeRef<ast::Expr>,
+
+    #[no_eq]
+    count: countme::Count<Expression<'static>>,
 }
 
 impl<'db> Expression<'db> {

crates/red_knot_python_semantic/src/semantic_index/symbol.rs

@@ -100,6 +100,9 @@ pub struct ScopeId<'db> {
     #[no_eq]
     #[return_ref]
     pub node: NodeWithScopeKind,
+
+    #[no_eq]
+    count: countme::Count<ScopeId<'static>>,
 }
 
 impl<'db> ScopeId<'db> {

crates/ruff_db/src/program.rs

@@ -65,7 +65,7 @@ impl std::fmt::Debug for TargetVersion {
 }
 
 /// Configures the search paths for module resolution.
-#[derive(Eq, PartialEq, Debug)]
+#[derive(Eq, PartialEq, Debug, Clone)]
 pub struct SearchPathSettings {
     /// List of user-provided paths that should take first priority in the module resolution.
     /// Examples in other type checkers are mypy's MYPYPATH environment variable,

crates/ruff_db/src/system.rs

@@ -9,7 +9,9 @@ use walk_directory::WalkDirectoryBuilder;
 
 use crate::file_revision::FileRevision;
 
-pub use self::path::{SystemPath, SystemPathBuf};
+pub use self::path::{
+    deduplicate_nested_paths, DeduplicatedNestedPathsIter, SystemPath, SystemPathBuf,
+};
 
 mod memory_fs;
 #[cfg(feature = "os")]

crates/ruff_db/src/system/path.rs

@@ -563,3 +563,60 @@ impl ruff_cache::CacheKey for SystemPathBuf {
         self.as_path().cache_key(hasher);
     }
 }
+
+/// Deduplicates identical paths and removes nested paths.
+///
+/// # Examples
+/// ```rust
+/// use ruff_db::system::{SystemPath, deduplicate_nested_paths};///
+///
+/// let paths = vec![SystemPath::new("/a/b/c"), SystemPath::new("/a/b"), SystemPath::new("/a/beta"), SystemPath::new("/a/b/c")];
+/// assert_eq!(deduplicate_nested_paths(paths).collect::<Vec<_>>(), &[SystemPath::new("/a/b"), SystemPath::new("/a/beta")]);
+/// ```
+pub fn deduplicate_nested_paths<'a, I>(paths: I) -> DeduplicatedNestedPathsIter<'a>
+where
+    I: IntoIterator<Item = &'a SystemPath>,
+{
+    DeduplicatedNestedPathsIter::new(paths)
+}
+
+pub struct DeduplicatedNestedPathsIter<'a> {
+    inner: std::vec::IntoIter<&'a SystemPath>,
+    next: Option<&'a SystemPath>,
+}
+
+impl<'a> DeduplicatedNestedPathsIter<'a> {
+    fn new<I>(paths: I) -> Self
+    where
+        I: IntoIterator<Item = &'a SystemPath>,
+    {
+        let mut paths = paths.into_iter().collect::<Vec<_>>();
+        // Sort the path to ensure that e.g. `/a/b/c`, comes right after `/a/b`.
+        paths.sort_unstable();
+
+        let mut iter = paths.into_iter();
+
+        Self {
+            next: iter.next(),
+            inner: iter,
+        }
+    }
+}
+
+impl<'a> Iterator for DeduplicatedNestedPathsIter<'a> {
+    type Item = &'a SystemPath;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let current = self.next.take()?;
+
+        for next in self.inner.by_ref() {
+            // Skip all paths that have the same prefix as the current path
+            if !next.starts_with(current) {
+                self.next = Some(next);
+                break;
+            }
+        }
+
+        Some(current)
+    }
+}