[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_dir<'a>(
+        &'a self,
+        path: &SystemPath,
+    ) -> Result<Box<dyn Iterator<Item = Result<DirEntry>> + 'a>>;
+
     fn as_any(&self) -> &dyn std::any::Any;
 }
 
@@ -78,7 +83,7 @@ impl Metadata {
     }
 }
 
-#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash)]
+#[derive(Copy, Clone, Eq, PartialEq, Debug, Hash, PartialOrd, Ord)]
 pub enum FileType {
     File,
     Directory,
@@ -98,3 +103,23 @@ impl FileType {
         matches!(self, FileType::Symlink)
     }
 }
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub struct DirEntry {
+    path: SystemPathBuf,
+    file_type: FileType,
+}
+
+impl DirEntry {
+    pub fn new(path: SystemPathBuf, file_type: FileType) -> Self {
+        Self { path, file_type }
+    }
+
+    pub fn path(&self) -> &SystemPath {
+        &self.path
+    }
+
+    pub fn file_type(&self) -> FileType {
+        self.file_type
+    }
+}

crates/ruff_db/src/system/memory_fs.rs

@@ -1,10 +1,11 @@
 use std::collections::BTreeMap;
+use std::iter::FusedIterator;
 use std::sync::{Arc, RwLock, RwLockWriteGuard};
 
 use camino::{Utf8Path, Utf8PathBuf};
 use filetime::FileTime;
 
-use crate::system::{FileType, Metadata, Result, SystemPath, SystemPathBuf};
+use crate::system::{DirEntry, FileType, Metadata, Result, SystemPath, SystemPathBuf};
 
 /// File system that stores all content in memory.
 ///
@@ -237,6 +238,24 @@ impl MemoryFileSystem {
         let normalized = SystemPath::absolute(path, &self.inner.cwd);
         normalized.into_utf8_path_buf()
     }
+
+    pub(crate) fn read_dir(&self, path: impl AsRef<SystemPath>) -> Result<DirectoryIterator> {
+        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());
+        };
+        // Is this inefficient? Yes, but this implementation of `System` is just for testing,
+        // and the lifetime issues get somewhat complicated without the cloning.
+        let contents = by_path.keys().cloned().collect();
+        Ok(DirectoryIterator {
+            paths: contents,
+            directory: normalized,
+            index: 0,
+            fs: self,
+        })
+    }
 }
 
 impl Default for MemoryFileSystem {
@@ -270,6 +289,40 @@ impl Entry {
     }
 }
 
+#[derive(Debug)]
+pub(crate) struct DirectoryIterator<'a> {
+    paths: Vec<Utf8PathBuf>,
+    directory: Utf8PathBuf,
+    index: usize,
+    fs: &'a MemoryFileSystem,
+}
+
+impl<'a> Iterator for DirectoryIterator<'a> {
+    type Item = Result<DirEntry>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            let path = self.paths.get(self.index)?;
+            self.index += 1;
+            let Some(parent) = path.parent() else {
+                continue;
+            };
+            if parent != self.directory {
+                continue;
+            }
+            let path = SystemPathBuf::from_utf8_path_buf(path.to_owned());
+            let file_type = if self.fs.is_directory(&path) {
+                FileType::Directory
+            } else {
+                FileType::File
+            };
+            break Some(Ok(DirEntry::new(path, file_type)));
+        }
+    }
+}
+
+impl FusedIterator for DirectoryIterator<'_> {}
+
 #[derive(Debug)]
 struct File {
     content: String,
@@ -349,7 +402,7 @@ mod tests {
     use std::io::ErrorKind;
     use std::time::Duration;
 
-    use crate::system::{MemoryFileSystem, Result, SystemPath};
+    use crate::system::{DirEntry, FileType, MemoryFileSystem, Result, SystemPath, SystemPathBuf};
 
     /// Creates a file system with the given files.
     ///
@@ -612,4 +665,31 @@ mod tests {
         let error = fs.remove_directory("a").unwrap_err();
         assert_eq!(error.kind(), ErrorKind::Other);
     }
+
+    #[test]
+    fn read_dir() {
+        let fs = with_files(["b.ts", "a/bar.py", "d.rs", "a/foo/bar.py", "a/baz.pyi"]);
+        let contents: Vec<DirEntry> = fs.read_dir("a").unwrap().map(Result::unwrap).collect();
+        let expected_contents = vec![
+            DirEntry::new(SystemPathBuf::from("/a/bar.py"), FileType::File),
+            DirEntry::new(SystemPathBuf::from("/a/baz.pyi"), FileType::File),
+            DirEntry::new(SystemPathBuf::from("/a/foo"), FileType::Directory),
+        ];
+        assert_eq!(contents, expected_contents)
+    }
+
+    #[test]
+    fn read_dir_nonexistent() {
+        let fs = MemoryFileSystem::new();
+        let error = fs.read_dir("doesnt_exist").unwrap_err();
+        assert_eq!(error.kind(), std::io::ErrorKind::NotFound);
+    }
+
+    #[test]
+    fn read_dir_on_file() {
+        let fs = with_files(["a.py"]);
+        let error = fs.read_dir("a.py").unwrap_err();
+        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,4 @@
-use crate::system::{FileType, Metadata, Result, System, SystemPath, SystemPathBuf};
+use crate::system::{DirEntry, FileType, Metadata, Result, System, SystemPath, SystemPathBuf};
 use filetime::FileTime;
 use std::any::Any;
 use std::sync::Arc;
@@ -68,6 +68,14 @@ impl System for OsSystem {
     fn as_any(&self) -> &dyn Any {
         self
     }
+
+    fn read_dir(&self, path: &SystemPath) -> Result<Box<dyn Iterator<Item = Result<DirEntry>>>> {
+        Ok(Box::new(
+            path.as_camino_path()
+                .read_dir_utf8()?
+                .map(|res| res.and_then(DirEntry::try_from)),
+        ))
+    }
 }
 
 impl From<std::fs::FileType> for FileType {
@@ -81,3 +89,74 @@ impl From<std::fs::FileType> for FileType {
         }
     }
 }
+
+impl TryFrom<camino::Utf8DirEntry> for DirEntry {
+    type Error = std::io::Error;
+
+    fn try_from(value: camino::Utf8DirEntry) -> Result<Self> {
+        Ok(Self {
+            path: SystemPathBuf::from_utf8_path_buf(value.path().to_path_buf()),
+            file_type: FileType::from(value.file_type()?),
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use tempfile::TempDir;
+
+    use super::*;
+
+    #[test]
+    fn read_dir() {
+        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<DirEntry> = fs
+            .read_dir(&tempdir_path.join("a"))
+            .unwrap()
+            .map(Result::unwrap)
+            .collect();
+        sorted_contents.sort();
+
+        let expected_contents = vec![
+            DirEntry::new(tempdir_path.join("a/bar.py"), FileType::File),
+            DirEntry::new(tempdir_path.join("a/baz.pyi"), FileType::File),
+            DirEntry::new(tempdir_path.join("a/foo"), FileType::Directory),
+        ];
+        assert_eq!(sorted_contents, expected_contents)
+    }
+
+    #[test]
+    fn read_dir_nonexistent() {
+        let fs = OsSystem::new("");
+        let result = fs.read_dir(SystemPath::new("doesnt_exist"));
+        assert!(result.is_err_and(|error| error.kind() == std::io::ErrorKind::NotFound));
+    }
+
+    #[test]
+    fn read_dir_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_dir(&tempdir_path.join("a.py"));
+        let Err(error) = result else {
+            panic!("Expected the read_dir() call to fail!");
+        };
+        // 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,5 @@
 use crate::files::File;
-use crate::system::{MemoryFileSystem, Metadata, OsSystem, System, SystemPath};
+use crate::system::{DirEntry, MemoryFileSystem, Metadata, OsSystem, Result, System, SystemPath};
 use crate::Db;
 use std::any::Any;
 
@@ -85,6 +85,16 @@ impl System for TestSystem {
     fn as_any(&self) -> &dyn Any {
         self
     }
+
+    fn read_dir<'a>(
+        &'a self,
+        path: &SystemPath,
+    ) -> Result<Box<dyn Iterator<Item = Result<DirEntry>> + 'a>> {
+        match &self.inner {
+            TestFileSystem::Os(fs) => fs.read_dir(path),
+            TestFileSystem::Stub(fs) => Ok(Box::new(fs.read_dir(path)?)),
+        }
+    }
 }
 
 /// Extension trait for databases that use [`TestSystem`].