Update PDF class, add CLI for fulltext

Author: breuleux

src/paperoni/__main__.py

@@ -2,7 +2,7 @@
 import json
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Annotated, Any
+from typing import Annotated, Any, Literal
 
 import yaml
 from gifnoc import add_overlay, cli
@@ -11,6 +11,8 @@
 
 from .config import config
 from .display import display, terminal_width
+from .fulltext.locate import locate_all
+from .fulltext.pdf import CachePolicies, get_pdf
 from .model import PaperInfo
 from .model.merge import merge_all
 from .refinement import fetch_all
@@ -50,7 +52,7 @@ class Discover:
     """Discover papers from various sources."""
 
     command: Annotated[
-        Any, FromEntryPoint("paperoni.discovery", wrap=lambda cls: Auto[cls().query])
+        Any, FromEntryPoint("paperoni.discovery", wrap=lambda cls: Auto[cls.query])
     ]
 
     # Output format
@@ -66,6 +68,35 @@ def run(self):
         self.format(papers)
 
 
+def locate(
+    # Reference to locate
+    # [positional]
+    ref: str,
+):
+    for url in locate_all(ref):
+        print(f"\033[36m[{url.info}]\033[0m {url.url}")
+
+
+def download(
+    # Reference to locate
+    # [positional]
+    # [nargs: +]
+    ref: list[str],
+    # Cache policy
+    # [alias: -p]
+    cache_policy: Literal["use", "use_best", "no_download", "force"] = "use",
+):
+    p = get_pdf(ref, cache_policy=getattr(CachePolicies, cache_policy.upper()))
+    print("Downloaded into:", p.pdf_path.resolve())
+
+
+@dataclass
+class Fulltext:
+    """Download and process fulltext."""
+
+    run: TaggedUnion[Auto[locate], Auto[download]]
+
+
 @dataclass
 class Refine:
     """Refine paper information."""
@@ -97,7 +128,7 @@ def run(self):
 class PaperoniInterface:
     """Paper database"""
 
-    command: TaggedUnion[Discover, Refine]
+    command: TaggedUnion[Discover, Refine, Fulltext]
 
     def run(self):
         self.command.run()

src/paperoni/discovery/openalex.py

@@ -300,6 +300,7 @@ def _reconstruct_abstract(cls, inverted: Dict[str, List[int]]) -> str:
 
 @dataclass
 class OpenAlex(Discoverer):
+    # Email associated with the query, for politeness
     mailto: str = field(default_factory=lambda: config.mailto)
 
     def query(

src/paperoni/fulltext/download.py

@@ -0,0 +1,121 @@
+import hashlib
+import shutil
+from dataclasses import dataclass
+
+from serieux import dump, load
+
+from ..config import config
+from .locate import URL, locate_all
+
+
+@dataclass
+class CachePolicy:
+    use: bool = True
+    download: bool = True
+    best: bool = False
+
+
+class CachePolicies:
+    USE = CachePolicy(use=True, download=True, best=False)
+    USE_BEST = CachePolicy(use=True, download=True, best=True)
+    NO_DOWNLOAD = CachePolicy(use=True, download=False, best=False)
+    FORCE = CachePolicy(use=False, download=True, best=True)
+
+
+@dataclass
+class ErrorData:
+    type: str
+    message: str
+
+
+@dataclass
+class Info:
+    id: str = None
+    title: str = None
+    ref: str = None
+
+
+@dataclass
+class PDF:
+    source: URL
+    hash: str = None
+    info: Info = None
+    success: bool = False
+    error: ErrorData = None
+
+    def __post_init__(self):
+        if self.hash is None:
+            self.hash = hashlib.sha256(str(self.source).encode()).hexdigest()
+        self.directory = config.data_path / "pdf" / self.hash
+        self.meta_path = self.directory / "meta.yaml"
+        self.pdf_path = self.directory / "fulltext.pdf"
+
+    def ensure(self):
+        self.directory.mkdir(parents=True, exist_ok=True)
+
+    def load(self):
+        if self.meta_path.exists():
+            return load(PDF, self.meta_path)
+        else:
+            return self
+
+    def dump(self):
+        self.ensure()
+        dump(PDF, self, dest=self.meta_path)
+
+    def clear(self):
+        assert self.directory.resolve().is_relative_to(config.data_path.resolve())
+        shutil.rmtree(self.directory, ignore_errors=True)
+
+    def fetch(self):
+        self.ensure()
+        try:
+            config.fetch.download(
+                url=self.source.url,
+                filename=self.pdf_path,
+            )
+            self.success = True
+            return self.pdf_path
+        except Exception as exc:
+            self.error = ErrorData(
+                type=type(exc).__name__,
+                message=str(exc),
+            )
+            self.success = False
+            raise
+        finally:
+            self.dump()
+
+    def fulltext(self, cache_policy: CachePolicy = CachePolicies.USE):
+        if not cache_policy.use or not self.success or not self.pdf_path.exists():
+            if not cache_policy.download:
+                raise Exception(
+                    f"No PDF for {self.source.url} and cache policy prevents downloading"
+                )
+            return self.fetch()
+        else:
+            return self.pdf_path
+
+
+def get_pdf(refs, cache_policy: CachePolicy = CachePolicies.USE):
+    if isinstance(refs, str):
+        refs = [refs]
+
+    if cache_policy.use and not cache_policy.best:
+        urls = [url for ref in refs for url in locate_all(ref)]
+        for url in urls:
+            if (p := PDF(url).load()).success:
+                return p
+
+    exceptions = []
+    for ref in refs:
+        for url in locate_all(ref):
+            p = PDF(url).load()
+            try:
+                p.fulltext(cache_policy=cache_policy)
+                return p
+            except Exception as exc:
+                exceptions.append(exc)
+                continue
+
+    raise ExceptionGroup("No fulltext found for any reference", exceptions)