use pdf_extract::content::{Content, Operation};
use pdf_extract::{Document, Object, Stream, dictionary};
use semlith::chunk;
use std::fs;
use std::path::{Path, PathBuf};
fn fixtures() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures")
}
fn pdf(pages: usize) -> Vec<u8> {
let mut doc = Document::with_version("1.5");
let pages_id = doc.new_object_id();
let font = doc.add_object(dictionary! {
"Type" => "Font", "Subtype" => "Type1", "BaseFont" => "Courier",
});
let resources = doc.add_object(dictionary! { "Font" => dictionary! { "F1" => font } });
let mut kids: Vec<Object> = Vec::new();
for n in 1..=pages {
let content = Content {
operations: vec![
Operation::new("BT", vec![]),
Operation::new("Tf", vec!["F1".into(), 24.into()]),
Operation::new("Td", vec![72.into(), 720.into()]),
Operation::new("Tj", vec![Object::string_literal(format!("page {n}"))]),
Operation::new("ET", vec![]),
],
};
let stream = doc.add_object(Stream::new(dictionary! {}, content.encode().unwrap()));
let page = doc.add_object(dictionary! {
"Type" => "Page", "Parent" => pages_id, "Contents" => stream,
});
kids.push(page.into());
}
doc.objects.insert(
pages_id,
Object::Dictionary(dictionary! {
"Type" => "Pages",
"Kids" => kids,
"Count" => pages as i64,
"Resources" => resources,
"MediaBox" => vec![0.into(), 0.into(), 612.into(), 792.into()],
}),
);
let catalog = doc.add_object(dictionary! { "Type" => "Catalog", "Pages" => pages_id });
doc.trailer.set("Root", catalog);
let mut out = Vec::new();
doc.save_to(&mut out).unwrap();
out
}
fn units_of(name: &str, bytes: &[u8]) -> Option<i64> {
chunk::extract_counted(Path::new(name), bytes)
.unwrap_or_else(|why| panic!("{name} was skipped: {}", why.as_str()))
.1
}
const CSV: &str = "name,amount,note\nalpha,1,\"a, quoted\"\nbeta,,\n";
#[test]
fn every_counted_format_counts_what_the_file_holds() {
let expected = [
("deck.pptx", 12),
("deck.odp", 3),
("sheet.xlsx", 18),
("sheet.ods", 8),
("analysis.ipynb", 3),
];
for (name, units) in expected {
let bytes = fs::read(fixtures().join(name)).unwrap();
assert_eq!(units_of(name, &bytes), Some(units), "{name}");
}
assert_eq!(units_of("book.pdf", &pdf(3)), Some(3));
assert_eq!(units_of("t.csv", CSV.as_bytes()), Some(7));
assert_eq!(units_of("t.tsv", b"a\tb\n\tc\n"), Some(3));
let docx = fs::read(fixtures().join("notes.docx")).unwrap();
assert_eq!(units_of("notes.docx", &docx), None);
assert_eq!(units_of("main.rs", b"fn main() {}\n"), None);
}
#[test]
fn counting_pages_leaves_a_pdfs_text_as_it_was() {
for pages in [1, 2, 7] {
let bytes = pdf(pages);
let before = pdf_extract::extract_text_from_mem(&bytes).unwrap();
let after = chunk::extract(Path::new("x.pdf"), &bytes).unwrap();
assert_eq!(after, before, "{pages} pages");
assert!(after.contains(&format!("page {pages}")));
}
}
#[test]
#[ignore = "downloads an embedding model on first run"]
fn an_older_store_gains_its_counts_without_re_embedding() {
let corpus = tempfile::tempdir().unwrap();
let store = tempfile::tempdir().unwrap();
for name in [
"deck.pptx",
"deck.odp",
"sheet.xlsx",
"sheet.ods",
"analysis.ipynb",
] {
fs::copy(fixtures().join(name), corpus.path().join(name)).unwrap();
}
fs::write(corpus.path().join("book.pdf"), pdf(4)).unwrap();
fs::write(corpus.path().join("table.csv"), CSV).unwrap();
fs::write(corpus.path().join("notes.md"), "# Notes\n\nplain prose\n").unwrap();
let index = |s: &mut semlith::Semlith| {
s.index_paths(&[corpus.path().to_path_buf()], |_, _| {})
.unwrap()
};
let mut s = semlith::Semlith::open(store.path(), None).unwrap();
s.quiet = true;
let first = index(&mut s);
assert_eq!(first.indexed, 8, "{first:?}");
let db = rusqlite::Connection::open(store.path().join("store.db")).unwrap();
let kinds = |db: &rusqlite::Connection| {
let corpus = semlith::store::corpus(db, |_| String::new()).unwrap();
let mut rows: Vec<(String, i64, i64, i64)> = corpus
.kinds
.into_iter()
.map(|k| (k.name, k.count, k.units, k.uncounted))
.collect();
rows.sort();
rows
};
let counted = vec![
("Notebook".to_string(), 1, 3, 0),
("PDF".to_string(), 1, 4, 0),
("Slide deck".to_string(), 2, 15, 0),
("Spreadsheet".to_string(), 3, 33, 0),
("Text and code".to_string(), 1, 0, 0),
];
assert_eq!(kinds(&db), counted);
let snapshot = |db: &rusqlite::Connection| -> (Vec<i64>, i64) {
let mut q = db.prepare("SELECT id FROM chunks ORDER BY id").unwrap();
let ids = q
.query_map([], |r| r.get(0))
.unwrap()
.collect::<Result<Vec<i64>, _>>()
.unwrap();
let fts: i64 = db
.query_row("SELECT COUNT(*) FROM chunks_fts", [], |r| r.get(0))
.unwrap();
(ids, fts)
};
let before = snapshot(&db);
db.execute("UPDATE files SET units = NULL", []).unwrap();
let waiting = kinds(&db);
assert!(
waiting
.iter()
.filter(|k| k.0 != "Text and code")
.all(|k| k.3 == k.1 && k.2 == 0),
"{waiting:?}"
);
s.plan(&[corpus.path().to_path_buf()]).unwrap();
let again = index(&mut s);
assert_eq!(again.indexed, 0, "a count re-embedded a file: {again:?}");
assert_eq!(kinds(&db), counted);
assert_eq!(snapshot(&db), before, "chunks or FTS rows moved");
}