doc_extractor 0.1.0

Extract text chunks from common document formats.
Documentation
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
      );
    }
  }
}