use std::collections::HashSet;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use pagesift_api::types::Document;
use super::{PageOut, page_url};
pub struct DirOut {
dir: PathBuf,
used: Mutex<HashSet<String>>,
}
impl DirOut {
pub fn create(dir: impl Into<PathBuf>) -> std::io::Result<Self> {
let dir = dir.into();
std::fs::create_dir_all(&dir)?;
Ok(Self {
dir,
used: Mutex::new(HashSet::new()),
})
}
pub fn dir(&self) -> &Path {
&self.dir
}
pub fn summary(&self, json: &[u8]) -> std::io::Result<()> {
std::fs::write(self.dir.join("summary.json"), json)
}
fn file_name(&self, url: &str) -> String {
let trimmed = url.split_once("://").map_or(url, |(_, rest)| rest);
let mut name: String = trimmed
.chars()
.map(|c| {
if c.is_ascii_alphanumeric() || c == '-' || c == '.' {
c
} else {
'_'
}
})
.collect();
name.truncate(150);
let name = name.trim_matches('_').to_string();
let mut used = self.used.lock().unwrap_or_else(|e| e.into_inner());
let mut candidate = format!("{name}.md");
let mut n = 2;
while !used.insert(candidate.clone()) {
candidate = format!("{name}-{n}.md");
n += 1;
}
candidate
}
}
pub(super) fn index_line(doc: &Document, url: &str, file: &str) -> serde_json::Value {
let mut line = serde_json::json!({ "url": url, "file": file, "title": doc.metadata.title });
if let Some(p) = doc.metadata.pagesift.as_ref() {
if let Some(r) = p.relevance {
line["relevance"] = serde_json::json!(r);
}
if let Some(t) = &p.crawl {
line["depth"] = serde_json::json!(t.depth);
line["parent"] = serde_json::json!(t.parent);
line["score"] = serde_json::json!(t.score);
if let Some(a) = &t.attachment_of {
line["attachmentOf"] = serde_json::json!(a);
}
}
}
if !doc.attachments.is_empty() {
line["images"] =
serde_json::json!(doc.attachments.iter().map(|a| &a.id).collect::<Vec<_>>());
}
if !doc.image_assets.is_empty() {
line["imageAssets"] = serde_json::json!(doc.image_assets);
}
line["status"] = serde_json::json!(doc.metadata.status_code);
if let Some(w) = &doc.warning {
line["warning"] = serde_json::json!(w);
}
line
}
impl DirOut {
fn write_images(&self, doc: &Document) -> std::io::Result<String> {
let mut md = doc.markdown.clone().unwrap_or_default();
if doc.attachments.is_empty() {
return Ok(md);
}
let images = self.dir.join("images");
std::fs::create_dir_all(&images)?;
for a in &doc.attachments {
if a.id.is_empty() || !a.id.chars().all(|c| c.is_ascii_alphanumeric() || c == '.') {
continue;
}
if let Some(data) = &a.data {
let path = images.join(&a.id);
if !path.exists() {
std::fs::write(&path, data)?;
}
}
if !a.inline_data
&& let Some(src) = &a.source
{
let local = format!("images/{}", a.id);
md = md
.replace(&format!("]({src})"), &format!("]({local})"))
.replace(&format!("]({src} \""), &format!("]({local} \""));
}
}
Ok(md)
}
}
impl PageOut for DirOut {
fn page(&self, doc: &Document) -> std::io::Result<()> {
let url = page_url(doc);
let name = self.file_name(&url);
let markdown = self.write_images(doc)?;
std::fs::write(self.dir.join(&name), markdown)?;
let mut index = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(self.dir.join("index.jsonl"))?;
writeln!(index, "{}", index_line(doc, &url, &name))?;
eprintln!("{url} -> {name}");
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use pagesift_api::types::Attachment;
#[test]
fn asset_modes_keep_links_and_include_outcomes_in_index() {
let dir = std::env::temp_dir().join(format!("pagesift-dir-assets-{}", std::process::id()));
let out = DirOut::create(&dir).unwrap();
let mut a = Attachment::new(b"image".to_vec(), "png", Some("https://e.x/a.png".into()));
a.inline_data = true;
let mut doc = Document {
markdown: Some(format!(" ", a.id)),
image_assets: vec![pagesift_api::types::ImageAssetResult {
source: "https://e.x/a.png".into(),
status: pagesift_api::types::ImageAssetStatus::Ok,
id: Some(a.id.clone()),
mime: Some("image/png".into()),
bytes: Some(a.bytes),
reason: None,
}],
attachments: vec![a.clone()],
..Default::default()
};
doc.metadata.source_url = Some("https://e.x/asset".into());
out.page(&doc).unwrap();
assert_eq!(
std::fs::read_to_string(dir.join("e.x_asset.md")).unwrap(),
doc.markdown.unwrap()
);
assert_eq!(
std::fs::read(dir.join("images").join(&a.id)).unwrap(),
b"image"
);
let index: serde_json::Value = serde_json::from_str(
std::fs::read_to_string(dir.join("index.jsonl"))
.unwrap()
.lines()
.last()
.unwrap(),
)
.unwrap();
assert_eq!(index["imageAssets"][0]["id"], a.id);
std::fs::remove_dir_all(dir).unwrap();
}
#[test]
fn images_are_written_once_and_links_become_local() {
let dir = std::env::temp_dir().join(format!("pagesift-dirout-{}", std::process::id()));
let out = DirOut::create(&dir).unwrap();
let web = Attachment::new(
vec![0x89, b'P', b'N', b'G'],
"png",
Some("https://e.x/a.png".into()),
);
let pdf = Attachment::new(vec![0xFF, 0xD8, 0xFF], "jpg", None);
let mut doc = Document {
markdown: Some(format!(
"\n![b](https://e.x/a.png \"t\")\n\n",
pdf.id
)),
attachments: vec![web.clone(), pdf.clone()],
..Default::default()
};
doc.metadata.source_url = Some("https://e.x/page".into());
out.page(&doc).unwrap();
let md = std::fs::read_to_string(dir.join("e.x_page.md")).unwrap();
assert_eq!(
md,
format!(
"\n![b](images/{0} \"t\")\n\n",
web.id, pdf.id
)
);
assert_eq!(
std::fs::read(dir.join("images").join(&web.id)).unwrap(),
[0x89, b'P', b'N', b'G']
);
let index = std::fs::read_to_string(dir.join("index.jsonl")).unwrap();
assert!(index.contains(&web.id));
std::fs::remove_dir_all(&dir).unwrap();
}
}