From 6632d632398442f15ab4800a7161a158b945c198 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Wed, 19 Aug 2026 16:47:37 +0200 Subject: [PATCH] refactor: rework document/directory resolution and add its unit tests - collect_documents now partitions --doc paths into files vs. directories up front and handles each with its own loop, instead of a single branching if/else per path; the directory branch resolves entries via a fallible iterator chain (? inside a Result-returning map, collected with Result::transpose) rather than a nested for loop - extracted chunk_document (path, splitter) -> Vec as a pure function, and Document::from_chunk(path, index, chunk, total) to own the chunk's source-string formatting, both previously inlined in a function that mutated a shared Vec in place - renamed push_document -> chunk_document to match what it now does (reads and chunks a file into owned Documents) instead of what it used to do (push one Document into a caller-supplied Vec) - added unit tests for the reworked logic: multiple files/directories in one call, a mix of both in one call, an unreadable file inside a scanned directory (skipped, siblings kept) vs. an unreadable directory itself (fails the whole call, unlike a file), and direct tests of chunk_document and Document::from_chunk in isolation Verified with cargo build/test (34 passed)/clippy -D warnings/fmt --check. --- src/documents.rs | 206 +++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 182 insertions(+), 24 deletions(-) diff --git a/src/documents.rs b/src/documents.rs index a167330..3b34919 100644 --- a/src/documents.rs +++ b/src/documents.rs @@ -20,6 +20,25 @@ pub(crate) struct Document { pub(crate) text: String, } +impl Document { + /// Builds a `Document` for the `index`-th of `total` chunks split from + /// `path`, labeling multi-chunk files with a "(part N/total)" suffix so + /// retrieved excerpts can be traced back to their position in the + /// source file. + fn from_chunk(path: &Path, index: usize, chunk: &str, total: usize) -> Self { + let source = if total > 1 { + format!("{} (part {}/{total})", path.display(), index + 1) + } else { + path.display().to_string() + }; + + Self { + source, + text: chunk.to_string(), + } + } +} + impl rig::Embed for Document { fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError> { embedder.embed(self.text.clone()); @@ -37,42 +56,48 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result, Vec<&PathBuf>) = + paths.iter().partition(|path| path.is_dir()); + + for dir in directories { + std::fs::read_dir(dir)? + .map(|entry| -> anyhow::Result> { let entry = entry?; - if entry.file_type()?.is_file() { - push_document(&entry.path(), &splitter, &mut documents); - } - } - } else { - push_document(path, &splitter, &mut documents); - } + Ok(entry.file_type()?.is_file().then(|| entry.path())) + }) + .filter_map(Result::transpose) + .collect::>>()? + .iter() + .for_each(|path| documents.extend(chunk_document(path, &splitter))); } + files + .iter() + .for_each(|file| documents.extend(chunk_document(file, &splitter))); + Ok(documents) } -fn push_document(path: &Path, splitter: &TextSplitter, documents: &mut Vec) { +/// Reads `path` and splits it into chunked `Document`s via `splitter`. An +/// unreadable path (permissions, non-UTF-8, gone by the time it's read) +/// yields an empty `Vec` and a warning rather than failing the caller. +fn chunk_document(path: &Path, splitter: &TextSplitter) -> Vec { match std::fs::read_to_string(path) { Ok(text) => { let chunks: Vec<&str> = splitter.chunks(&text).collect(); - let total = chunks.len(); - let source = |i: usize| { - if total > 1 { - format!("{} (part {}/{total})", path.display(), i + 1) - } else { - path.display().to_string() - } - }; + let chunk_count = chunks.len(); - documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document { - source: source(i), - text: chunk.to_string(), - })); + chunks + .into_iter() + .enumerate() + .map(|(chunk_index, chunk)| { + Document::from_chunk(path, chunk_index, chunk, chunk_count) + }) + .collect() } Err(err) => { - tracing::warn!(path = %path.display(), %err, "skipping unreadable document") + tracing::warn!(path = %path.display(), %err, "skipping unreadable document"); + Vec::new() } } } @@ -160,4 +185,137 @@ mod tests { vec!["the text to embed".to_string()] ); } + + #[test] + fn collect_documents_reads_multiple_explicit_files() { + let dir = tempdir().unwrap(); + let a = dir.path().join("a.txt"); + let b = dir.path().join("b.txt"); + std::fs::write(&a, "a").unwrap(); + std::fs::write(&b, "b").unwrap(); + + let mut documents = collect_documents(&[a, b]).unwrap(); + documents.sort_by(|x, y| x.text.cmp(&y.text)); + + assert_eq!(documents.len(), 2); + assert_eq!(documents[0].text, "a"); + assert_eq!(documents[1].text, "b"); + } + + #[test] + fn collect_documents_reads_multiple_directories() { + let first = tempdir().unwrap(); + let second = tempdir().unwrap(); + std::fs::write(first.path().join("a.txt"), "a").unwrap(); + std::fs::write(second.path().join("b.txt"), "b").unwrap(); + + let mut documents = + collect_documents(&[first.path().to_path_buf(), second.path().to_path_buf()]).unwrap(); + documents.sort_by(|x, y| x.text.cmp(&y.text)); + + assert_eq!(documents.len(), 2); + assert_eq!(documents[0].text, "a"); + assert_eq!(documents[1].text, "b"); + } + + #[test] + fn collect_documents_handles_a_mix_of_files_and_directories_in_one_call() { + let explicit = tempdir().unwrap(); + let explicit_file = explicit.path().join("explicit.txt"); + std::fs::write(&explicit_file, "explicit").unwrap(); + + let scanned = tempdir().unwrap(); + std::fs::write(scanned.path().join("scanned.txt"), "scanned").unwrap(); + + let mut documents = + collect_documents(&[explicit_file, scanned.path().to_path_buf()]).unwrap(); + documents.sort_by(|x, y| x.text.cmp(&y.text)); + + assert_eq!(documents.len(), 2); + assert_eq!(documents[0].text, "explicit"); + assert_eq!(documents[1].text, "scanned"); + } + + #[cfg(unix)] + #[test] + fn collect_documents_skips_an_unreadable_file_inside_a_directory_but_keeps_the_rest() { + use std::os::unix::fs::PermissionsExt; + + let dir = tempdir().unwrap(); + let readable = dir.path().join("readable.txt"); + let unreadable = dir.path().join("unreadable.txt"); + std::fs::write(&readable, "readable").unwrap(); + std::fs::write(&unreadable, "unreadable").unwrap(); + std::fs::set_permissions(&unreadable, std::fs::Permissions::from_mode(0o000)).unwrap(); + + let documents = collect_documents(&[dir.path().to_path_buf()]); + + std::fs::set_permissions(&unreadable, std::fs::Permissions::from_mode(0o644)).unwrap(); + + let documents = documents.unwrap(); + assert_eq!(documents.len(), 1); + assert_eq!(documents[0].text, "readable"); + } + + #[cfg(unix)] + #[test] + fn collect_documents_fails_outright_on_an_unreadable_directory() { + use std::os::unix::fs::PermissionsExt; + + let dir = tempdir().unwrap(); + std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o000)).unwrap(); + + let result = collect_documents(&[dir.path().to_path_buf()]); + + std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o755)).unwrap(); + + assert!( + result.is_err(), + "an unreadable directory itself should fail the whole call, unlike an unreadable file" + ); + } + + #[test] + fn chunk_document_splits_a_readable_file_by_the_given_splitter() { + let dir = tempdir().unwrap(); + let file = dir.path().join("small.txt"); + std::fs::write(&file, "hello").unwrap(); + let splitter = TextSplitter::new(CHUNK_CHARS); + + let chunks = chunk_document(&file, &splitter); + + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].text, "hello"); + assert_eq!(chunks[0].source, file.display().to_string()); + } + + #[test] + fn chunk_document_returns_empty_for_an_unreadable_path() { + let dir = tempdir().unwrap(); + let missing = dir.path().join("does-not-exist.txt"); + let splitter = TextSplitter::new(CHUNK_CHARS); + + assert!(chunk_document(&missing, &splitter).is_empty()); + } + + #[test] + fn document_from_chunk_uses_the_bare_path_when_there_is_only_one_chunk() { + let path = Path::new("/tmp/notes.txt"); + + let document = Document::from_chunk(path, 0, "text", 1); + + assert_eq!(document.source, "/tmp/notes.txt"); + assert_eq!(document.text, "text"); + } + + #[test] + fn document_from_chunk_appends_a_part_suffix_when_there_are_multiple_chunks() { + let path = Path::new("/tmp/notes.txt"); + + let first = Document::from_chunk(path, 0, "a", 3); + let last = Document::from_chunk(path, 2, "c", 3); + + assert_eq!(first.source, "/tmp/notes.txt (part 1/3)"); + assert_eq!(last.source, "/tmp/notes.txt (part 3/3)"); + } }