[red-knot] Add a read_directory() method to the ruff_db::system::System trait

crates/ruff_db/Cargo.toml

@@ -27,3 +27,4 @@ zip = { workspace = true }
 
 [dev-dependencies]
 insta = { workspace = true }
+tempfile = { workspace = true }

crates/ruff_db/src/system.rs

@@ -54,6 +54,11 @@ pub trait System {
     /// Returns the current working directory
     fn current_directory(&self) -> &SystemPath;
 
+    fn read_directory<'a>(
+        &'a self,
+        path: &SystemPath,
+    ) -> Result<Box<dyn Iterator<Item = Result<DirectoryEntry>> + 'a>>;
+
     fn as_any(&self) -> &dyn std::any::Any;
 }
 
@@ -98,3 +103,29 @@ impl FileType {
         matches!(self, FileType::Symlink)
     }
 }
+
+#[derive(Debug)]
+pub struct DirectoryEntry {
+    path: SystemPathBuf,
+    file_type: Result<FileType>,
+}
+
+impl DirectoryEntry {
+    pub fn new(path: SystemPathBuf, file_type: Result<FileType>) -> Self {
+        Self { path, file_type }
+    }
+
+    pub fn path(&self) -> &SystemPath {
+        &self.path
+    }
+
+    pub fn file_type(&self) -> &Result<FileType> {
+        &self.file_type
+    }
+}
+
+impl PartialEq for DirectoryEntry {
+    fn eq(&self, other: &Self) -> bool {
+        self.path == other.path
+    }
+}

crates/ruff_db/src/system/memory_fs.rs

@@ -4,7 +4,7 @@ use std::sync::{Arc, RwLock, RwLockWriteGuard};
 use camino::{Utf8Path, Utf8PathBuf};
 use filetime::FileTime;
 
-use crate::system::{FileType, Metadata, Result, SystemPath, SystemPathBuf};
+use crate::system::{DirectoryEntry, FileType, Metadata, Result, SystemPath, SystemPathBuf};
 
 /// File system that stores all content in memory.
 ///
@@ -237,6 +237,34 @@ impl MemoryFileSystem {
         let normalized = SystemPath::absolute(path, &self.inner.cwd);
         normalized.into_utf8_path_buf()
     }
+
+    pub(crate) fn read_directory(
+        &self,
+        path: impl AsRef<SystemPath>,
+    ) -> Result<impl Iterator<Item = Result<DirectoryEntry>> + '_> {
+        let by_path = self.inner.by_path.read().unwrap();
+        let normalized = self.normalize_path(path.as_ref());
+        let entry = by_path.get(&normalized).ok_or_else(not_found)?;
+        if entry.is_file() {
+            return Err(not_a_directory());
+        };
+        Ok(by_path
+            .range(normalized.clone()..)
+            .skip(1)
+            .take_while(|(path, _)| path.starts_with(&normalized))
+            .filter_map(|(path, entry)| {
+                if path.parent()? == normalized {
+                    Some(Ok(DirectoryEntry {
+                        path: SystemPathBuf::from_utf8_path_buf(path.to_owned()),
+                        file_type: Ok(entry.file_type()),
+                    }))
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>()
+            .into_iter())
+    }
 }
 
 impl Default for MemoryFileSystem {
@@ -268,6 +296,13 @@ impl Entry {
     const fn is_file(&self) -> bool {
         matches!(self, Entry::File(_))
     }
+
+    const fn file_type(&self) -> FileType {
+        match self {
+            Self::File(_) => FileType::File,
+            Self::Directory(_) => FileType::Directory,
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -349,7 +384,9 @@ mod tests {
     use std::io::ErrorKind;
     use std::time::Duration;
 
-    use crate::system::{MemoryFileSystem, Result, SystemPath};
+    use crate::system::{
+        DirectoryEntry, FileType, MemoryFileSystem, Result, SystemPath, SystemPathBuf,
+    };
 
     /// Creates a file system with the given files.
     ///
@@ -612,4 +649,39 @@ mod tests {
         let error = fs.remove_directory("a").unwrap_err();
         assert_eq!(error.kind(), ErrorKind::Other);
     }
+
+    #[test]
+    fn read_directory() {
+        let fs = with_files(["b.ts", "a/bar.py", "d.rs", "a/foo/bar.py", "a/baz.pyi"]);
+        let contents: Vec<DirectoryEntry> = fs
+            .read_directory("a")
+            .unwrap()
+            .map(Result::unwrap)
+            .collect();
+        let expected_contents = vec![
+            DirectoryEntry::new(SystemPathBuf::from("/a/bar.py"), Ok(FileType::File)),
+            DirectoryEntry::new(SystemPathBuf::from("/a/baz.pyi"), Ok(FileType::File)),
+            DirectoryEntry::new(SystemPathBuf::from("/a/foo"), Ok(FileType::Directory)),
+        ];
+        assert_eq!(contents, expected_contents)
+    }
+
+    #[test]
+    fn read_directory_nonexistent() {
+        let fs = MemoryFileSystem::new();
+        let Err(error) = fs.read_directory("doesnt_exist") else {
+            panic!("Expected this to fail");
+        };
+        assert_eq!(error.kind(), std::io::ErrorKind::NotFound);
+    }
+
+    #[test]
+    fn read_directory_on_file() {
+        let fs = with_files(["a.py"]);
+        let Err(error) = fs.read_directory("a.py") else {
+            panic!("Expected this to fail");
+        };
+        assert_eq!(error.kind(), std::io::ErrorKind::Other);
+        assert!(error.to_string().contains("Not a directory"));
+    }
 }

crates/ruff_db/src/system/os.rs

@@ -1,4 +1,6 @@
-use crate::system::{FileType, Metadata, Result, System, SystemPath, SystemPathBuf};
+use crate::system::{
+    DirectoryEntry, FileType, Metadata, Result, System, SystemPath, SystemPathBuf,
+};
 use filetime::FileTime;
 use std::any::Any;
 use std::sync::Arc;
@@ -68,6 +70,17 @@ impl System for OsSystem {
     fn as_any(&self) -> &dyn Any {
         self
     }
+
+    fn read_directory(
+        &self,
+        path: &SystemPath,
+    ) -> Result<Box<dyn Iterator<Item = Result<DirectoryEntry>>>> {
+        Ok(Box::new(
+            path.as_camino_path()
+                .read_dir_utf8()?
+                .map(|res| res.map(DirectoryEntry::from)),
+        ))
+    }
 }
 
 impl From<std::fs::FileType> for FileType {
@@ -81,3 +94,74 @@ impl From<std::fs::FileType> for FileType {
         }
     }
 }
+
+impl From<camino::Utf8DirEntry> for DirectoryEntry {
+    fn from(value: camino::Utf8DirEntry) -> Self {
+        let file_type = value.file_type().map(FileType::from);
+        Self {
+            path: SystemPathBuf::from_utf8_path_buf(value.into_path()),
+            file_type,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use tempfile::TempDir;
+
+    use super::*;
+
+    #[test]
+    fn read_directory() {
+        let tempdir = TempDir::new().unwrap();
+        let tempdir_path = tempdir.path();
+        std::fs::create_dir_all(tempdir_path.join("a/foo")).unwrap();
+        let files = &["b.ts", "a/bar.py", "d.rs", "a/foo/bar.py", "a/baz.pyi"];
+        for path in files {
+            std::fs::File::create(tempdir_path.join(path)).unwrap();
+        }
+
+        let tempdir_path = SystemPath::from_std_path(tempdir_path).unwrap();
+        let fs = OsSystem::new(tempdir_path);
+
+        let mut sorted_contents: Vec<DirectoryEntry> = fs
+            .read_directory(&tempdir_path.join("a"))
+            .unwrap()
+            .map(Result::unwrap)
+            .collect();
+        sorted_contents.sort_by(|a, b| a.path.cmp(&b.path));
+
+        let expected_contents = vec![
+            DirectoryEntry::new(tempdir_path.join("a/bar.py"), Ok(FileType::File)),
+            DirectoryEntry::new(tempdir_path.join("a/baz.pyi"), Ok(FileType::File)),
+            DirectoryEntry::new(tempdir_path.join("a/foo"), Ok(FileType::Directory)),
+        ];
+        assert_eq!(sorted_contents, expected_contents)
+    }
+
+    #[test]
+    fn read_directory_nonexistent() {
+        let fs = OsSystem::new("");
+        let result = fs.read_directory(SystemPath::new("doesnt_exist"));
+        assert!(result.is_err_and(|error| error.kind() == std::io::ErrorKind::NotFound));
+    }
+
+    #[test]
+    fn read_directory_on_file() {
+        let tempdir = TempDir::new().unwrap();
+        let tempdir_path = tempdir.path();
+        std::fs::File::create(tempdir_path.join("a.py")).unwrap();
+
+        let tempdir_path = SystemPath::from_std_path(tempdir_path).unwrap();
+        let fs = OsSystem::new(tempdir_path);
+        let result = fs.read_directory(&tempdir_path.join("a.py"));
+        let Err(error) = result else {
+            panic!("Expected the read_dir() call to fail!");
+        };
+        dbg!(error.to_string());
+        // We can't assert the error kind here because it's apparently an unstable feature!
+        // https://github.com/rust-lang/rust/issues/86442
+        // assert_eq!(error.kind(), std::io::ErrorKind::NotADirectory);
+        assert!(error.to_string().contains("Not a directory"));
+    }
+}

crates/ruff_db/src/system/path.rs

@@ -303,6 +303,12 @@ impl SystemPath {
         self.0.as_std_path()
     }
 
+    /// Returns the [`Utf8Path`] for the file.
+    #[inline]
+    pub fn as_camino_path(&self) -> &Utf8Path {
+        &self.0
+    }
+
     pub fn from_std_path(path: &Path) -> Option<&SystemPath> {
         Some(SystemPath::new(Utf8Path::from_path(path)?))
     }

crates/ruff_db/src/system/test.rs

@@ -1,5 +1,7 @@
 use crate::files::File;
-use crate::system::{MemoryFileSystem, Metadata, OsSystem, System, SystemPath};
+use crate::system::{
+    DirectoryEntry, MemoryFileSystem, Metadata, OsSystem, Result, System, SystemPath,
+};
 use crate::Db;
 use std::any::Any;
 
@@ -85,6 +87,16 @@ impl System for TestSystem {
     fn as_any(&self) -> &dyn Any {
         self
     }
+
+    fn read_directory<'a>(
+        &'a self,
+        path: &SystemPath,
+    ) -> Result<Box<dyn Iterator<Item = Result<DirectoryEntry>> + 'a>> {
+        match &self.inner {
+            TestFileSystem::Os(fs) => fs.read_directory(path),
+            TestFileSystem::Stub(fs) => Ok(Box::new(fs.read_directory(path)?)),
+        }
+    }
 }
 
 /// Extension trait for databases that use [`TestSystem`].