use office_oxide::{Document as OfficeDocument, DocumentFormat};
use super::*;
pub struct OfficeLoader {
document: OfficeDocument,
}
impl OfficeLoader {
pub fn new<R: Read + Seek + Send + 'static>(reader: R, format: DocumentFormat) -> Result<Self, LoaderError> {
Ok(Self {
document: OfficeDocument::from_reader(reader, format)?,
})
}
fn extract_text(&self) -> String {
self.document.to_markdown()
}
fn extract_text_to_doc(&self) -> Document {
Document::new(self.extract_text())
}
}
impl Loader for OfficeLoader {
fn load(self) -> LoaderResult<Vec<Document>> {
let doc = self.extract_text_to_doc();
Ok(vec![doc])
}
}
#[cfg(test)]
mod tests {
use std::{fs::read, io::Cursor, path::PathBuf};
use office_oxide::{DocumentFormat, create::create_from_markdown_to_writer};
use crate::Doc;
use super::*;
fn get_fixtures_path() -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../fixtures")
}
#[test]
fn test_parse_docx() {
let docx_buffer = include_bytes!("../../../../fixtures/demo.docx");
let parsed_buffer = include_str!("../../../../fixtures/demo.docx.md");
{
let loader = OfficeLoader::new(Cursor::new(docx_buffer), DocumentFormat::Docx).unwrap();
let documents = loader.load().unwrap();
assert_eq!(documents.len(), 1);
assert_eq!(documents[0].page_content, parsed_buffer);
}
{
let loader = OfficeLoader::new(Cursor::new(docx_buffer), DocumentFormat::Docx).unwrap();
let documents = loader.load_and_split(TokenSplitter::default()).unwrap();
for (idx, doc) in documents.into_iter().enumerate() {
assert_eq!(
doc.page_content,
String::from_utf8_lossy(&read(get_fixtures_path().join(format!("demo.docx.{}.md", idx))).unwrap())
);
}
}
}
#[test]
fn test_parse_generated_ooxml_formats() {
for format in [DocumentFormat::Docx, DocumentFormat::Xlsx, DocumentFormat::Pptx] {
let mut buffer = Cursor::new(Vec::new());
create_from_markdown_to_writer("# Office loader smoke\n\nShared parser text.", format, &mut buffer).unwrap();
let loader = OfficeLoader::new(Cursor::new(buffer.into_inner()), format).unwrap();
let documents = loader.load().unwrap();
assert_eq!(documents.len(), 1);
assert!(
documents[0].page_content.contains("Office loader smoke"),
"{format:?} output was: {}",
documents[0].page_content
);
assert!(
documents[0].page_content.contains("Shared parser text"),
"{format:?} output was: {}",
documents[0].page_content
);
}
}
#[test]
fn test_auto_load_generated_ooxml_formats() {
for (file_name, format) in [
("generated.docx", DocumentFormat::Docx),
("generated.xlsx", DocumentFormat::Xlsx),
("generated.pptx", DocumentFormat::Pptx),
] {
let mut buffer = Cursor::new(Vec::new());
create_from_markdown_to_writer("# Auto Office loader\n\nLoaded through Doc.", format, &mut buffer).unwrap();
let doc = Doc::new(file_name, &buffer.into_inner()).unwrap();
assert_eq!(doc.chunks.len(), 1);
assert!(
doc.chunks[0].content.contains("Auto Office loader"),
"{format:?} output was: {}",
doc.chunks[0].content
);
assert!(
doc.chunks[0].content.contains("Loaded through Doc"),
"{format:?} output was: {}",
doc.chunks[0].content
);
}
}
}