mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
566 lines
20 KiB
Python
566 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Zotero Corpus Pipeline
|
|
======================
|
|
Reads your Zotero SQLite library, scans local PDF directories for arXiv IDs,
|
|
fetches missing metadata from arXiv API, and reconciles everything into a
|
|
unified local index.
|
|
|
|
Usage:
|
|
python zotero_corpus_pipeline.py --scan /home/allaun/Downloads/data
|
|
python zotero_corpus_pipeline.py --report
|
|
python zotero_corpus_pipeline.py --reconcile
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import sqlite3
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
from collections import defaultdict
|
|
from dataclasses import asdict, dataclass, field
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Set
|
|
|
|
# ── Configuration ────────────────────────────────────────────────────────────
|
|
ZOTERO_DB = Path.home() / "Zotero" / "zotero.sqlite"
|
|
DEFAULT_SCAN_ROOTS = [
|
|
Path.home() / "Downloads" / "data" / "Downloads_from_internet",
|
|
Path.home() / "Downloads" / "data" / "literature",
|
|
Path.home() / "Zotero" / "storage",
|
|
]
|
|
INDEX_DB = Path.home() / "Research Stack" / "data" / "substrate_index.db"
|
|
|
|
# ── Data Classes ─────────────────────────────────────────────────────────────
|
|
@dataclass
|
|
class ZoteroItem:
|
|
item_id: int
|
|
key: str
|
|
item_type: str
|
|
title: Optional[str] = None
|
|
date: Optional[str] = None
|
|
doi: Optional[str] = None
|
|
url: Optional[str] = None
|
|
extra: Optional[str] = None
|
|
creators: List[str] = field(default_factory=list)
|
|
collections: List[str] = field(default_factory=list)
|
|
|
|
@dataclass
|
|
class LocalPDF:
|
|
path: Path
|
|
arxiv_id: Optional[str] = None
|
|
title_guess: Optional[str] = None
|
|
doi_guess: Optional[str] = None
|
|
file_size: int = 0
|
|
|
|
@dataclass
|
|
class ArxivMeta:
|
|
arxiv_id: str
|
|
title: str
|
|
authors: List[str]
|
|
summary: str
|
|
published: str
|
|
updated: str
|
|
primary_category: str
|
|
categories: List[str]
|
|
pdf_url: str
|
|
|
|
# ── Zotero Reader ────────────────────────────────────────────────────────────
|
|
class ZoteroReader:
|
|
def __init__(self, db_path: Path = ZOTERO_DB):
|
|
self.db_path = db_path
|
|
self.conn: Optional[sqlite3.Connection] = None
|
|
|
|
def connect(self) -> "ZoteroReader":
|
|
if not self.db_path.exists():
|
|
raise FileNotFoundError(f"Zotero DB not found: {self.db_path}")
|
|
self.conn = sqlite3.connect(f"file:{self.db_path}?mode=ro", uri=True)
|
|
return self
|
|
|
|
def close(self):
|
|
if self.conn:
|
|
self.conn.close()
|
|
self.conn = None
|
|
|
|
def _get_field_map(self) -> Dict[int, str]:
|
|
assert self.conn
|
|
cur = self.conn.execute("SELECT fieldID, fieldName FROM fields")
|
|
return {row[0]: row[1] for row in cur.fetchall()}
|
|
|
|
def _get_collection_map(self) -> Dict[int, str]:
|
|
assert self.conn
|
|
cur = self.conn.execute("SELECT collectionID, collectionName FROM collections")
|
|
return {row[0]: row[1] for row in cur.fetchall()}
|
|
|
|
def get_items(self, limit: Optional[int] = None) -> List[ZoteroItem]:
|
|
assert self.conn
|
|
field_map = self._get_field_map()
|
|
collection_map = self._get_collection_map()
|
|
|
|
# Pull all item metadata in one go
|
|
query = """
|
|
SELECT i.itemID, it.typeName, i.key, i.itemTypeID
|
|
FROM items i
|
|
JOIN itemTypes it ON i.itemTypeID = it.itemTypeID
|
|
WHERE i.itemTypeID NOT IN (1, 14)
|
|
ORDER BY i.dateAdded DESC
|
|
"""
|
|
if limit:
|
|
query += f" LIMIT {limit}"
|
|
|
|
cur = self.conn.execute(query)
|
|
items: Dict[int, ZoteroItem] = {}
|
|
for row in cur.fetchall():
|
|
item_id, type_name, key, _ = row
|
|
items[item_id] = ZoteroItem(
|
|
item_id=item_id, key=key, item_type=type_name
|
|
)
|
|
|
|
# Metadata values
|
|
meta_cur = self.conn.execute(
|
|
"""
|
|
SELECT id.itemID, id.fieldID, v.value
|
|
FROM itemData id
|
|
JOIN itemDataValues v ON id.valueID = v.valueID
|
|
WHERE id.itemID IN ({})
|
|
""".format(",".join(map(str, items.keys())))
|
|
)
|
|
for item_id, field_id, value in meta_cur.fetchall():
|
|
field_name = field_map.get(field_id, "")
|
|
item = items[item_id]
|
|
if field_name == "title":
|
|
item.title = value
|
|
elif field_name == "date":
|
|
item.date = value
|
|
elif field_name == "DOI":
|
|
item.doi = value
|
|
elif field_name == "url":
|
|
item.url = value
|
|
elif field_name == "extra":
|
|
item.extra = value
|
|
|
|
# Creators
|
|
creator_cur = self.conn.execute(
|
|
"""
|
|
SELECT ic.itemID, c.firstName, c.lastName
|
|
FROM itemCreators ic
|
|
JOIN creators c ON ic.creatorID = c.creatorID
|
|
WHERE ic.itemID IN ({})
|
|
ORDER BY ic.orderIndex
|
|
""".format(",".join(map(str, items.keys())))
|
|
)
|
|
for item_id, first, last in creator_cur.fetchall():
|
|
name = f"{first or ''} {last or ''}".strip()
|
|
if name:
|
|
items[item_id].creators.append(name)
|
|
|
|
# Collections
|
|
col_cur = self.conn.execute(
|
|
"""
|
|
SELECT ci.itemID, ci.collectionID
|
|
FROM collectionItems ci
|
|
WHERE ci.itemID IN ({})
|
|
""".format(",".join(map(str, items.keys())))
|
|
)
|
|
for item_id, col_id in col_cur.fetchall():
|
|
col_name = collection_map.get(col_id)
|
|
if col_name:
|
|
items[item_id].collections.append(col_name)
|
|
|
|
return list(items.values())
|
|
|
|
def get_stats(self) -> Dict:
|
|
assert self.conn
|
|
stats = {}
|
|
cur = self.conn.execute("SELECT COUNT(*) FROM items WHERE itemTypeID NOT IN (1, 14)")
|
|
stats["total_items"] = cur.fetchone()[0]
|
|
cur = self.conn.execute(
|
|
"""SELECT it.typeName, COUNT(*) FROM items i
|
|
JOIN itemTypes it ON i.itemTypeID = it.itemTypeID
|
|
WHERE i.itemTypeID NOT IN (1, 14)
|
|
GROUP BY it.typeName"""
|
|
)
|
|
stats["by_type"] = {row[0]: row[1] for row in cur.fetchall()}
|
|
cur = self.conn.execute("SELECT COUNT(*) FROM itemAttachments")
|
|
stats["attachments"] = cur.fetchone()[0]
|
|
cur = self.conn.execute("SELECT COUNT(*) FROM collections")
|
|
stats["collections"] = cur.fetchone()[0]
|
|
return stats
|
|
|
|
# ── PDF Scanner ──────────────────────────────────────────────────────────────
|
|
class PDFScanner:
|
|
ARXIV_RE = re.compile(r"(\d{4}\.\d{4,5}(?:v\d+)?)")
|
|
ARXIV_FN_RE = re.compile(r"(?:ar[xX]iv[_-]?)?(\d{4}\.\d{4,5}(?:v\d+)?)")
|
|
DOI_RE = re.compile(r"10\.\d{4,9}/[-._;()/:A-Z0-9a-z]+")
|
|
|
|
def __init__(self, roots: List[Path]):
|
|
self.roots = roots
|
|
|
|
def scan(self) -> List[LocalPDF]:
|
|
pdfs: List[LocalPDF] = []
|
|
for root in self.roots:
|
|
if not root.exists():
|
|
print(f"[skip] Missing root: {root}")
|
|
continue
|
|
for path in root.rglob("*.pdf"):
|
|
if not path.is_file():
|
|
continue
|
|
pdf = self._analyze(path)
|
|
pdfs.append(pdf)
|
|
return pdfs
|
|
|
|
def _analyze(self, path: Path) -> LocalPDF:
|
|
name = path.stem
|
|
# Try filename patterns first
|
|
arxiv_match = self.ARXIV_FN_RE.search(name)
|
|
arxiv_id = arxiv_match.group(1) if arxiv_match else None
|
|
|
|
# Fallback: scan first 8KB of PDF for DOI/arXiv
|
|
doi_guess = None
|
|
if not arxiv_id:
|
|
try:
|
|
with open(path, "rb") as f:
|
|
header = f.read(8192).decode("utf-8", errors="ignore")
|
|
doi_match = self.DOI_RE.search(header)
|
|
if doi_match:
|
|
doi_guess = doi_match.group(0)
|
|
am = self.ARXIV_RE.search(header)
|
|
if am:
|
|
arxiv_id = am.group(1)
|
|
except Exception:
|
|
pass
|
|
|
|
title_guess = self._clean_title(name)
|
|
return LocalPDF(
|
|
path=path,
|
|
arxiv_id=arxiv_id,
|
|
title_guess=title_guess,
|
|
doi_guess=doi_guess,
|
|
file_size=path.stat().st_size,
|
|
)
|
|
|
|
@staticmethod
|
|
def _clean_title(filename: str) -> str:
|
|
t = filename.replace("_", " ").replace("-", " ")
|
|
t = re.sub(r"\b\d{4}\.\d{4,5}v?\d*\b", "", t) # strip arxiv ids
|
|
t = re.sub(r"\s+", " ", t).strip()
|
|
return t
|
|
|
|
# ── arXiv API ────────────────────────────────────────────────────────────────
|
|
class ArxivClient:
|
|
BASE = "http://export.arxiv.org/api/query"
|
|
|
|
def fetch(self, arxiv_id: str) -> Optional[ArxivMeta]:
|
|
# strip vN suffix for API query
|
|
clean_id = re.sub(r"v\d+$", "", arxiv_id)
|
|
url = f"{self.BASE}?id_list={clean_id}&max_results=1"
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=30) as resp:
|
|
xml = resp.read().decode("utf-8")
|
|
return self._parse(xml, arxiv_id)
|
|
except Exception as e:
|
|
print(f"[arxiv] Failed to fetch {arxiv_id}: {e}")
|
|
return None
|
|
|
|
def _parse(self, xml: str, original_id: str) -> Optional[ArxivMeta]:
|
|
import xml.etree.ElementTree as ET
|
|
ns = {"atom": "http://www.w3.org/2005/Atom",
|
|
"arxiv": "http://arxiv.org/schemas/atom"}
|
|
root = ET.fromstring(xml)
|
|
entry = root.find("atom:entry", ns)
|
|
if entry is None:
|
|
return None
|
|
def tag(t): return entry.find(f"atom:{t}", ns)
|
|
title = (tag("title") or entry.find("title")).text.strip()
|
|
summary = (tag("summary") or entry.find("summary")).text.strip()
|
|
published = (tag("published") or entry.find("published")).text.strip()
|
|
updated_el = tag("updated") or entry.find("updated")
|
|
updated = updated_el.text.strip() if updated_el is not None else published
|
|
authors = [a.find("atom:name", ns).text for a in entry.findall("atom:author", ns)]
|
|
cat_el = entry.find("arxiv:primary_category", ns)
|
|
primary = cat_el.attrib.get("term", "") if cat_el is not None else ""
|
|
cats = [c.attrib.get("term", "") for c in entry.findall("atom:category", ns)]
|
|
pdf_url = f"https://arxiv.org/pdf/{original_id}.pdf"
|
|
return ArxivMeta(
|
|
arxiv_id=original_id,
|
|
title=title,
|
|
authors=authors,
|
|
summary=summary,
|
|
published=published,
|
|
updated=updated,
|
|
primary_category=primary,
|
|
categories=cats,
|
|
pdf_url=pdf_url,
|
|
)
|
|
|
|
# ── Index Manager ────────────────────────────────────────────────────────────
|
|
class IndexManager:
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS zotero_items (
|
|
zotero_key TEXT PRIMARY KEY,
|
|
item_id INTEGER,
|
|
item_type TEXT,
|
|
title TEXT,
|
|
date TEXT,
|
|
doi TEXT,
|
|
url TEXT,
|
|
extra TEXT,
|
|
creators TEXT, -- JSON list
|
|
collections TEXT -- JSON list
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS local_pdfs (
|
|
path TEXT PRIMARY KEY,
|
|
arxiv_id TEXT,
|
|
title_guess TEXT,
|
|
doi_guess TEXT,
|
|
file_size INTEGER,
|
|
zotero_key TEXT,
|
|
metadata_fetched INTEGER DEFAULT 0
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS arxiv_meta (
|
|
arxiv_id TEXT PRIMARY KEY,
|
|
title TEXT,
|
|
authors TEXT, -- JSON list
|
|
summary TEXT,
|
|
published TEXT,
|
|
updated TEXT,
|
|
primary_category TEXT,
|
|
categories TEXT, -- JSON list
|
|
pdf_url TEXT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS reconciler_log (
|
|
run_time TEXT,
|
|
zotero_count INTEGER,
|
|
pdf_count INTEGER,
|
|
matched_count INTEGER,
|
|
unmatched_pdf_count INTEGER,
|
|
notes TEXT
|
|
);
|
|
"""
|
|
|
|
def __init__(self, db_path: Path = INDEX_DB):
|
|
self.db_path = db_path
|
|
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.conn = sqlite3.connect(str(self.db_path))
|
|
self.conn.executescript(self.SCHEMA)
|
|
self.conn.commit()
|
|
|
|
def close(self):
|
|
self.conn.close()
|
|
|
|
def sync_zotero(self, items: List[ZoteroItem]):
|
|
with self.conn:
|
|
self.conn.execute("DELETE FROM zotero_items")
|
|
for it in items:
|
|
self.conn.execute(
|
|
"""INSERT INTO zotero_items VALUES (?,?,?,?,?,?,?,?,?,?)""",
|
|
(
|
|
it.key,
|
|
it.item_id,
|
|
it.item_type,
|
|
it.title,
|
|
it.date,
|
|
it.doi,
|
|
it.url,
|
|
it.extra,
|
|
json.dumps(it.creators),
|
|
json.dumps(it.collections),
|
|
),
|
|
)
|
|
|
|
def sync_pdfs(self, pdfs: List[LocalPDF]):
|
|
seen: Set[str] = set()
|
|
with self.conn:
|
|
self.conn.execute("DELETE FROM local_pdfs")
|
|
for p in pdfs:
|
|
path_str = str(p.path)
|
|
if path_str in seen:
|
|
continue
|
|
seen.add(path_str)
|
|
self.conn.execute(
|
|
"""INSERT INTO local_pdfs (path, arxiv_id, title_guess, doi_guess, file_size)
|
|
VALUES (?,?,?,?,?)""",
|
|
(path_str, p.arxiv_id, p.title_guess, p.doi_guess, p.file_size),
|
|
)
|
|
|
|
def store_arxiv_meta(self, meta: ArxivMeta):
|
|
with self.conn:
|
|
self.conn.execute(
|
|
"""INSERT OR REPLACE INTO arxiv_meta VALUES (?,?,?,?,?,?,?,?,?,?)""",
|
|
(
|
|
meta.arxiv_id,
|
|
meta.title,
|
|
json.dumps(meta.authors),
|
|
meta.summary,
|
|
meta.published,
|
|
meta.updated,
|
|
meta.primary_category,
|
|
json.dumps(meta.categories),
|
|
meta.pdf_url,
|
|
),
|
|
)
|
|
|
|
def reconcile(self) -> Dict:
|
|
"""Match local PDFs to Zotero items by DOI or arXiv ID."""
|
|
with self.conn:
|
|
# Match by DOI
|
|
self.conn.execute(
|
|
"""UPDATE local_pdfs
|
|
SET zotero_key = (
|
|
SELECT zotero_key FROM zotero_items
|
|
WHERE zotero_items.doi = local_pdfs.doi_guess
|
|
LIMIT 1
|
|
)
|
|
WHERE doi_guess IS NOT NULL"""
|
|
)
|
|
# Match by arXiv ID in URL or extra
|
|
self.conn.execute(
|
|
"""UPDATE local_pdfs
|
|
SET zotero_key = (
|
|
SELECT z.zotero_key FROM zotero_items z
|
|
WHERE (z.url LIKE '%' || local_pdfs.arxiv_id || '%'
|
|
OR z.extra LIKE '%' || local_pdfs.arxiv_id || '%')
|
|
LIMIT 1
|
|
)
|
|
WHERE arxiv_id IS NOT NULL AND zotero_key IS NULL"""
|
|
)
|
|
# Match by title similarity (naive)
|
|
self.conn.execute(
|
|
"""UPDATE local_pdfs
|
|
SET zotero_key = (
|
|
SELECT z.zotero_key FROM zotero_items z
|
|
WHERE z.title IS NOT NULL
|
|
AND local_pdfs.title_guess IS NOT NULL
|
|
AND lower(trim(z.title)) = lower(trim(local_pdfs.title_guess))
|
|
LIMIT 1
|
|
)
|
|
WHERE zotero_key IS NULL"""
|
|
)
|
|
|
|
cur = self.conn.execute("SELECT COUNT(*) FROM zotero_items")
|
|
z_count = cur.fetchone()[0]
|
|
cur = self.conn.execute("SELECT COUNT(*) FROM local_pdfs")
|
|
p_count = cur.fetchone()[0]
|
|
cur = self.conn.execute(
|
|
"SELECT COUNT(DISTINCT zotero_key) FROM local_pdfs WHERE zotero_key IS NOT NULL"
|
|
)
|
|
matched = cur.fetchone()[0]
|
|
cur = self.conn.execute(
|
|
"SELECT COUNT(*) FROM local_pdfs WHERE zotero_key IS NULL"
|
|
)
|
|
unmatched = cur.fetchone()[0]
|
|
return {
|
|
"zotero_items": z_count,
|
|
"local_pdfs": p_count,
|
|
"matched_pdfs": matched,
|
|
"unmatched_pdfs": unmatched,
|
|
}
|
|
|
|
def report(self) -> str:
|
|
lines = []
|
|
cur = self.conn.execute("SELECT * FROM zotero_items LIMIT 10")
|
|
lines.append("== Sample Zotero Items ==")
|
|
for row in cur.fetchall():
|
|
lines.append(f" {row[0]} | {row[3]} | {row[8]}")
|
|
|
|
cur = self.conn.execute(
|
|
"SELECT path, arxiv_id, zotero_key FROM local_pdfs LIMIT 10"
|
|
)
|
|
lines.append("\n== Sample Local PDFs ==")
|
|
for row in cur.fetchall():
|
|
lines.append(f" {Path(row[0]).name} | arXiv:{row[1]} | Zotero:{row[2]}")
|
|
|
|
cur = self.conn.execute(
|
|
"SELECT arxiv_id, title FROM arxiv_meta LIMIT 10"
|
|
)
|
|
lines.append("\n== Cached arXiv Metadata ==")
|
|
for row in cur.fetchall():
|
|
lines.append(f" {row[0]} | {row[1][:60]}")
|
|
|
|
cur = self.conn.execute(
|
|
"""SELECT zotero_key, COUNT(*) FROM local_pdfs
|
|
WHERE zotero_key IS NOT NULL GROUP BY zotero_key ORDER BY COUNT(*) DESC LIMIT 5"""
|
|
)
|
|
lines.append("\n== PDFs per Zotero Item (top 5) ==")
|
|
for row in cur.fetchall():
|
|
lines.append(f" {row[0]}: {row[1]} PDFs")
|
|
|
|
return "\n".join(lines)
|
|
|
|
# ── Main ─────────────────────────────────────────────────────────────────────
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Zotero Corpus Pipeline")
|
|
parser.add_argument("--scan", nargs="*", help="Extra directories to scan for PDFs")
|
|
parser.add_argument("--report", action="store_true", help="Print index report")
|
|
parser.add_argument("--reconcile", action="store_true", help="Run reconciler")
|
|
parser.add_argument("--fetch-arxiv", action="store_true", help="Fetch arXiv metadata for unmatched PDFs")
|
|
parser.add_argument("--stats", action="store_true", help="Print Zotero stats")
|
|
args = parser.parse_args()
|
|
|
|
# 1. Read Zotero
|
|
print("[1/4] Connecting to Zotero...")
|
|
zr = ZoteroReader().connect()
|
|
if args.stats:
|
|
stats = zr.get_stats()
|
|
print(json.dumps(stats, indent=2))
|
|
zr.close()
|
|
return
|
|
|
|
items = zr.get_items()
|
|
print(f" -> {len(items)} Zotero items loaded")
|
|
|
|
# 2. Scan PDFs
|
|
print("[2/4] Scanning local PDF corpus...")
|
|
roots = list(DEFAULT_SCAN_ROOTS)
|
|
if args.scan:
|
|
roots += [Path(p) for p in args.scan]
|
|
scanner = PDFScanner(roots)
|
|
pdfs = scanner.scan()
|
|
print(f" -> {len(pdfs)} PDFs found")
|
|
|
|
# 3. Update index
|
|
print("[3/4] Writing index...")
|
|
idx = IndexManager()
|
|
idx.sync_zotero(items)
|
|
idx.sync_pdfs(pdfs)
|
|
|
|
# 4. Reconcile
|
|
if args.reconcile:
|
|
print("[4/4] Reconciling...")
|
|
rec = idx.reconcile()
|
|
print(f" -> {rec['matched_pdfs']} matched, {rec['unmatched_pdfs']} unmatched")
|
|
|
|
# 5. Fetch arXiv metadata for unmatched PDFs with arxiv_id
|
|
if args.fetch_arxiv:
|
|
print("[5/5] Fetching arXiv metadata...")
|
|
client = ArxivClient()
|
|
cur = idx.conn.execute(
|
|
"""SELECT arxiv_id FROM local_pdfs
|
|
WHERE arxiv_id IS NOT NULL
|
|
AND zotero_key IS NULL
|
|
AND arxiv_id NOT IN (SELECT arxiv_id FROM arxiv_meta)"""
|
|
)
|
|
missing = [row[0] for row in cur.fetchall()]
|
|
print(f" -> {len(missing)} missing metadata records")
|
|
for i, aid in enumerate(missing, 1):
|
|
meta = client.fetch(aid)
|
|
if meta:
|
|
idx.store_arxiv_meta(meta)
|
|
if i % 10 == 0:
|
|
print(f" ...{i}/{len(missing)} fetched")
|
|
time.sleep(3) # be polite to arXiv
|
|
|
|
if args.report:
|
|
print(idx.report())
|
|
|
|
idx.close()
|
|
zr.close()
|
|
print("Done.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|