[ty] Fix panic 'missing root' when handling completion request (#20917)

Author: MichaReiser

crates/ruff_db/src/files.rs

@@ -193,6 +193,17 @@ impl Files {
         roots.at(&absolute)
     }
 
+    /// The same as [`Self::root`] but panics if no root is found.
+    #[track_caller]
+    pub fn expect_root(&self, db: &dyn Db, path: &SystemPath) -> FileRoot {
+        if let Some(root) = self.root(db, path) {
+            return root;
+        }
+
+        let roots = self.inner.roots.read().unwrap();
+        panic!("No root found for path '{path}'. Known roots: {roots:#?}");
+    }
+
     /// Adds a new root for `path` and returns the root.
     ///
     /// The root isn't added nor is the file root's kind updated if a root for `path` already exists.

crates/ruff_db/src/files/file_root.rs

@@ -81,6 +81,8 @@ impl FileRoots {
             }
         }
 
+        tracing::debug!("Adding new file root '{path}' of kind {kind:?}");
+
         // normalize the path to use `/` separators and escape the '{' and '}' characters,
         // which matchit uses for routing parameters
         let mut route = normalized_path.replace('{', "{{").replace('}', "}}");

crates/ty_python_semantic/src/module_resolver/list.rs

@@ -65,17 +65,15 @@ fn list_modules_in<'db>(
     db: &'db dyn Db,
     search_path: SearchPathIngredient<'db>,
 ) -> Vec<Module<'db>> {
+    tracing::debug!("Listing modules in search path '{}'", search_path.path(db));
     let mut lister = Lister::new(db, search_path.path(db));
     match search_path.path(db).as_path() {
         SystemOrVendoredPathRef::System(system_search_path) => {
             // Read the revision on the corresponding file root to
             // register an explicit dependency on this directory. When
             // the revision gets bumped, the cache that Salsa creates
             // for this routine will be invalidated.
-            let root = db
-                .files()
-                .root(db, system_search_path)
-                .expect("System search path should have a registered root");
+            let root = db.files().expect_root(db, system_search_path);
             let _ = root.revision(db);
 
             let Ok(it) = db.system().read_directory(system_search_path) else {
@@ -969,10 +967,6 @@ mod tests {
         std::os::unix::fs::symlink(foo.as_std_path(), bar.as_std_path())?;
 
         db.files().try_add_root(&db, &src, FileRootKind::Project);
-        db.files()
-            .try_add_root(&db, &site_packages, FileRootKind::LibrarySearchPath);
-        db.files()
-            .try_add_root(&db, &custom_typeshed, FileRootKind::LibrarySearchPath);
 
         Program::from_settings(
             &db,
@@ -1468,6 +1462,55 @@ not_a_directory
         );
     }
 
+    #[test]
+    fn editable_installs_into_first_party_search_path() {
+        let mut db = TestDb::new();
+
+        let src = SystemPath::new("/src");
+        let venv_site_packages = SystemPathBuf::from("/venv-site-packages");
+        let site_packages_pth = venv_site_packages.join("foo.pth");
+        let editable_install_location = src.join("x/y/a.py");
+
+        db.write_files([
+            (&site_packages_pth, "/src/x/y/"),
+            (&editable_install_location, ""),
+        ])
+        .unwrap();
+
+        db.files()
+            .try_add_root(&db, SystemPath::new("/src"), FileRootKind::Project);
+
+        Program::from_settings(
+            &db,
+            ProgramSettings {
+                python_version: PythonVersionWithSource::default(),
+                python_platform: PythonPlatform::default(),
+                search_paths: SearchPathSettings {
+                    site_packages_paths: vec![venv_site_packages],
+                    ..SearchPathSettings::new(vec![src.to_path_buf()])
+                }
+                .to_search_paths(db.system(), db.vendored())
+                .expect("Valid search path settings"),
+            },
+        );
+
+        insta::assert_debug_snapshot!(
+            list_snapshot_filter(&db, |m| m.name(&db).as_str() == "a"),
+            @r#"
+        [
+            Module::File("a", "editable", "/src/x/y/a.py", Module, None),
+        ]
+        "#,
+        );
+
+        let editable_root = db
+            .files()
+            .root(&db, &editable_install_location)
+            .expect("file root for editable install");
+
+        assert_eq!(editable_root.path(&db), src);
+    }
+
     #[test]
     fn multiple_site_packages_with_editables() {
         let mut db = TestDb::new();
@@ -1490,12 +1533,6 @@ not_a_directory
 
         db.files()
             .try_add_root(&db, SystemPath::new("/src"), FileRootKind::Project);
-        db.files()
-            .try_add_root(&db, &venv_site_packages, FileRootKind::LibrarySearchPath);
-        db.files()
-            .try_add_root(&db, &system_site_packages, FileRootKind::LibrarySearchPath);
-        db.files()
-            .try_add_root(&db, SystemPath::new("/x"), FileRootKind::LibrarySearchPath);
 
         Program::from_settings(
             &db,
@@ -1625,8 +1662,6 @@ not_a_directory
 
         db.files()
             .try_add_root(&db, &project_directory, FileRootKind::Project);
-        db.files()
-            .try_add_root(&db, &site_packages, FileRootKind::LibrarySearchPath);
 
         Program::from_settings(
             &db,

crates/ty_python_semantic/src/module_resolver/module.rs

@@ -175,10 +175,7 @@ fn all_submodule_names_for_package<'db>(
             // tree. When the revision gets bumped, the cache
             // that Salsa creates does for this routine will be
             // invalidated.
-            let root = db
-                .files()
-                .root(db, parent_directory)
-                .expect("System search path should have a registered root");
+            let root = db.files().expect_root(db, parent_directory);
             let _ = root.revision(db);
 
             db.system()

crates/ty_python_semantic/src/module_resolver/resolver.rs

@@ -348,9 +348,15 @@ impl SearchPaths {
         })
     }
 
+    /// Registers the file roots for all non-dynamically discovered search paths that aren't first-party.
     pub(crate) fn try_register_static_roots(&self, db: &dyn Db) {
         let files = db.files();
-        for path in self.static_paths.iter().chain(self.site_packages.iter()) {
+        for path in self
+            .static_paths
+            .iter()
+            .chain(self.site_packages.iter())
+            .chain(&self.stdlib_path)
+        {
             if let Some(system_path) = path.as_system_path() {
                 if !path.is_first_party() {
                     files.try_add_root(db, system_path, FileRootKind::LibrarySearchPath);
@@ -451,9 +457,7 @@ pub(crate) fn dynamic_resolution_paths<'db>(
             continue;
         }
 
-        let site_packages_root = files
-            .root(db, &site_packages_dir)
-            .expect("Site-package root to have been created");
+        let site_packages_root = files.expect_root(db, &site_packages_dir);
 
         // This query needs to be re-executed each time a `.pth` file
         // is added, modified or removed from the `site-packages` directory.
@@ -500,6 +504,23 @@ pub(crate) fn dynamic_resolution_paths<'db>(
                             "Adding editable installation to module resolution path {path}",
                             path = installation
                         );
+
+                        // Register a file root for editable installs that are outside any other root
+                        // (Most importantly, don't register a root for editable installations from the project
+                        // directory as that would change the durability of files within those folders).
+                        // Not having an exact file root for editable installs just means that
+                        // some queries (like `list_modules_in`) will run slightly more frequently
+                        // than they would otherwise.
+                        if let Some(dynamic_path) = search_path.as_system_path() {
+                            if files.root(db, dynamic_path).is_none() {
+                                files.try_add_root(
+                                    db,
+                                    dynamic_path,
+                                    FileRootKind::LibrarySearchPath,
+                                );
+                            }
+                        }
+
                         dynamic_paths.push(search_path);
                     }