Skills › Scientific Agent Skills › paper-lookup › _common.py
_common.py
227 lines · skills/paper-lookup/scripts/_common.py at 5f2bc03
1#!/usr/bin/env python3
2"""Shared helpers for the paper-lookup scripts. Standard library only.
4Three concerns are factored out here because all four CLIs need them and getting
5any of them subtly wrong is how a literature retrieval turns into a plausible
6lie:
8`read_input` / `emit`
9 Bounded stdin-or-path reading and JSON writing, so a 400 MB full-text pull
10 cannot exhaust memory unnoticed.
12`collapse_ws` / `strip_control`
13 API payloads are third-party text. Titles and abstracts arrive hard-wrapped,
14 and full text can carry control characters that corrupt a terminal or a
15 downstream parse.
17`Reconciliation`
18 Expected total versus retrieved total, in one place, because every paginated
19 API in this skill counts differently and the whole point is to fail visibly
20 when they disagree.
21"""
23from __future__ import annotations
25import json
26import re
27import sys
28from dataclasses import dataclass, field
29from pathlib import Path
30from typing import Any
32#: A single response should never be this large. PMC full text runs ~1 MB; a
33#: 64 MB payload means a bulk dump was piped in by mistake.
34MAX_INPUT_BYTES = 64 * 1024 * 1024
36_CONTROL = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
37_WHITESPACE = re.compile(r"\s+")
40class InputError(Exception):
41 """Bad input from the caller: unreadable path, oversized payload, bad JSON."""
44def read_input(source: str, *, max_bytes: int = MAX_INPUT_BYTES) -> str:
45 """Read `source`, or stdin when it is `-`, refusing anything oversized.
47 stdin is read in chunks rather than whole so that an accidental
48 `cat huge.xml | script` fails fast instead of after filling memory.
49 """
50 if source == "-":
51 chunks: list[bytes] = []
52 total = 0
53 stream = sys.stdin.buffer
54 while True:
55 chunk = stream.read(1024 * 1024)
56 if not chunk:
57 break
58 total += len(chunk)
59 if total > max_bytes:
60 raise InputError(
61 f"stdin exceeded {max_bytes} bytes; write it to a file and pass a path, "
62 "or slice the payload first"
63 )
64 chunks.append(chunk)
65 raw = b"".join(chunks)
66 else:
67 path = Path(source)
68 if not path.is_file():
69 raise InputError(f"not a file: {source}")
70 size = path.stat().st_size
71 if size > max_bytes:
72 raise InputError(f"{source} is {size} bytes, over the {max_bytes} byte limit")
73 raw = path.read_bytes()
75 if not raw.strip():
76 raise InputError("input was empty")
77 return raw.decode("utf-8", errors="replace")
80def load_json(source: str, *, max_bytes: int = MAX_INPUT_BYTES) -> Any:
81 """`read_input` plus a JSON parse, with the failure attributed to the source."""
82 text = read_input(source, max_bytes=max_bytes)
83 try:
84 return json.loads(text)
85 except json.JSONDecodeError as error:
86 where = "stdin" if source == "-" else source
87 raise InputError(f"{where} is not valid JSON: {error}") from error
90def emit(payload: Any, destination: str | None = None) -> None:
91 """Write `payload` as UTF-8 JSON to a path, or to stdout when None."""
92 text = json.dumps(payload, indent=2, ensure_ascii=False, sort_keys=False)
93 if destination is None:
94 sys.stdout.write(text + "\n")
95 else:
96 Path(destination).write_text(text + "\n", encoding="utf-8")
99def strip_control(text: str) -> str:
100 """Drop control characters, keeping tab, newline, and carriage return."""
101 return _CONTROL.sub("", text)
104def collapse_ws(text: str | None) -> str:
105 """Collapse all whitespace runs to single spaces and trim.
107 arXiv hard-wraps `<title>` and `<summary>` mid-sentence, and JATS indents
108 element text, so raw values compare unequal to the same string from any other
109 source. Every field this skill emits for display goes through here.
110 """
111 if not text:
112 return ""
113 return _WHITESPACE.sub(" ", strip_control(text)).strip()
116@dataclass
117class Reconciliation:
118 """Expected versus retrieved, and *why* they differ when they do.
120 Three outcomes, deliberately not collapsed into one boolean:
122 `complete`
123 The API said it was done and the counts agree. Nothing to caveat.
125 `stopped_at_limit`
126 A `--max-records` / `--max-calls` bound was reached. The result is
127 partial **because the caller asked for a partial result** -- honest, and
128 not an error. It still must be reported as partial, since presenting 100
129 of 697,030 as "the papers on X" is the misleading case this skill exists
130 to prevent.
132 shortfall
133 The walk believed it had finished, yet retrieved fewer than the total the
134 API reported. Records went missing. This is the one that must never pass
135 quietly -- it is what an out-of-step bioRxiv cursor produces.
137 `expected` is None for the several endpoints here that report no total at all
138 (bioRxiv DOI lookups, `/details/{N}`). That is a documented state, not a
139 failure.
140 """
142 expected: int | None = None
143 retrieved: int = 0
144 pages: int = 0
145 stopped_at_limit: bool = False
146 notes: list[str] = field(default_factory=list)
148 @property
149 def complete(self) -> bool:
150 """Did the walk retrieve everything the API said exists?"""
151 if self.stopped_at_limit:
152 return False
153 if self.expected is None:
154 return True
155 return self.retrieved == self.expected
157 @property
158 def ok(self) -> bool:
159 """Is the shortfall explained? False only when records went missing."""
160 return self.complete or self.stopped_at_limit
162 def note(self, message: str) -> None:
163 self.notes.append(message)
165 def as_dict(self) -> dict[str, Any]:
166 summary: dict[str, Any] = {
167 "expected_total": self.expected,
168 "retrieved_total": self.retrieved,
169 "pages_fetched": self.pages,
170 "complete": self.complete,
171 "stopped_at_limit": self.stopped_at_limit,
172 }
173 if self.expected is None:
174 summary["expected_total_note"] = (
175 "endpoint reports no total; retrieved_total is all that can be asserted"
176 )
177 elif self.retrieved != self.expected:
178 summary["shortfall"] = self.expected - self.retrieved
179 summary["shortfall_reason"] = (
180 "bounded by --max-records/--max-calls; raise the bound to continue"
181 if self.stopped_at_limit
182 else "UNEXPLAINED: the walk ended on its own but came up short -- records are missing"
183 )
184 if self.notes:
185 summary["notes"] = list(self.notes)
186 return summary
189#: Query parameters that must never appear in emitted provenance.
190#:
191#: Several of these APIs authenticate by query string rather than header, so the
192#: URL that was actually fetched contains the credential. Provenance is supposed
193#: to let someone repeat the call with *their own* key -- printing yours is a leak,
194#: not reproducibility. `email`/`mailto` are contact details rather than secrets,
195#: but they are still the caller's personal address and do not belong in output
196#: that gets pasted into a report.
197REDACTED_PARAMS = frozenset({"api_key", "apikey", "key", "email", "mailto", "tool"})
198REDACTION = "REDACTED"
201def redact_url(url: str) -> str:
202 """Replace credential query-parameter values with a placeholder.
204 The parameter *names* survive so the call stays reproducible: a reader can
205 see that `api_key` was supplied and substitute their own.
206 """
207 from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
209 parts = urlsplit(url)
210 if not parts.query:
211 return url
212 pairs = [
213 (name, REDACTION if name.lower() in REDACTED_PARAMS else value)
214 for name, value in parse_qsl(parts.query, keep_blank_values=True)
215 ]
216 return urlunsplit(parts._replace(query=urlencode(pairs)))
219def fail(message: str, code: int = 1) -> None:
220 """Write `message` to stderr and exit non-zero.
222 Scripts here exit non-zero on silent-failure conditions -- a JATS document
223 with no `<body>`, an arXiv Error entry, a pagination shortfall -- precisely
224 because the APIs return HTTP 200 for them.
225 """
226 sys.stderr.write(f"error: {message}\n")
227 raise SystemExit(code)