seroverse.client¶

class Client:
    """High-performance asynchronous HTTP client powered by `aiohttp` and `msgspec`.

    Args:
        base_url: Base endpoint URL.
        api_key_header: Optional HTTP header key for API authorization.
        api_key: Optional API key string.
        user_agent: User-Agent header value.
        timeout: Total request timeout in seconds.
        max_concurrency: Maximum concurrent in-flight requests per event loop.
        polite_delay: Optional minimum delay (in seconds) after live network requests.
        recorder: Optional `cassetter` recorder instance.
    """

    def __init__(
        self,
        base_url: str,
        api_key_header: str | None = None,
        api_key: str | None = None,
        user_agent: str = "seroverse-client/0.1",
        timeout: float = 30.0,
        *,
        max_concurrency: int = 5,
        polite_delay: float = 0.0,
        recorder: Any | None = None,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.api_key_header = api_key_header
        self._api_key: str | None = api_key
        self.user_agent = user_agent
        self.timeout = aiohttp.ClientTimeout(total=timeout)
        self.max_concurrency = max(1, int(max_concurrency))
        self.polite_delay = float(polite_delay)
        self.recorder = recorder
        self._session: aiohttp.ClientSession | None = None
        self._sem: asyncio.Semaphore | None = None
        self._sem_loop: asyncio.AbstractEventLoop | None = None

    @property
    def api_key(self) -> str | None:
        """Return the configured API key string, or `None`."""
        return self._api_key

    @staticmethod
    def encode_msgpack(obj: Any) -> bytes:
        """Serialize an object or dataclass tree into compact MessagePack bytes."""
        return msgspec.msgpack.encode(obj)

    @staticmethod
    def decode_msgpack[T](data: bytes, type_: type[T]) -> T:
        """Deserialize MessagePack bytes directly into a typed `@dataclass(frozen=True, slots=True)` structure."""
        return msgspec.msgpack.decode(data, type=type_)

    @staticmethod
    def save_msgpack(path: Path, obj: Any) -> None:
        """Atomically persist `obj` to `path` as binary MessagePack."""
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_name(f"{path.name}.{os.getpid()}.tmp")
        tmp.write_bytes(msgspec.msgpack.encode(obj))
        os.replace(tmp, path)

    @staticmethod
    def load_msgpack[T](path: Path, type_: type[T]) -> T:
        """Load and decode a binary MessagePack file into `type_`."""
        return msgspec.msgpack.decode(path.read_bytes(), type=type_)

    def __repr__(self) -> str:
        key_repr = "'**********'" if self._api_key else "None"
        return f"{self.__class__.__name__}(base_url='{self.base_url}', api_key={key_repr})"

    def __cache_fingerprint__(self) -> tuple[str, str, str]:
        """Returns a deterministic, credential-safe fingerprint tuple for content-addressable caches."""
        key_digest = hashlib.sha256((self._api_key or "").encode("utf-8")).hexdigest()[:16] if self._api_key else ""
        return (type(self).__qualname__, self.base_url, key_digest)

    def _headers(self) -> dict[str, str]:
        headers = {"User-Agent": self.user_agent, "Accept": "application/json"}
        if self.api_key_header and self._api_key:
            headers[self.api_key_header] = self._api_key
        return headers

    def _get_semaphore(self) -> asyncio.Semaphore:
        loop = asyncio.get_running_loop()
        if self._sem is None or self._sem_loop is not loop:
            self._sem = asyncio.Semaphore(self.max_concurrency)
            self._sem_loop = loop
        return self._sem

    async def polite_sleep(self, t_req: float, delay: float | None = None) -> None:
        """Sleep politely after a live network request, skipping sleep on fast cassette replays."""
        wait_s = self.polite_delay if delay is None else delay
        if wait_s <= 0.0:
            return
        elapsed = time.perf_counter() - t_req
        if not (self.recorder is not None and elapsed < 0.25) and elapsed > 0.02:
            await asyncio.sleep(wait_s)

    async def get_session(self) -> aiohttp.ClientSession:
        """Retrieve or instantiate the underlying `aiohttp.ClientSession` for the active event loop."""
        current_loop = asyncio.get_running_loop()
        if (
            self._session is None
            or self._session.closed
            or getattr(self._session, "_loop", None) is not current_loop
            or getattr(self._session, "_loop", current_loop).is_closed()
        ):
            self._session = aiohttp.ClientSession(
                headers=self._headers(),
                timeout=self.timeout,
            )
        return self._session

    async def __aenter__(self) -> Self:
        await self.get_session()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        await self.close()

    def __enter__(self) -> Self:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.close_sync()

    async def close(self) -> None:
        """Close the underlying `aiohttp.ClientSession` if open."""
        if self._session is not None:
            if not self._session.closed:
                loop = getattr(self._session, "_loop", None)
                if loop is None or not loop.is_closed():
                    await self._session.close()
            self._session = None

    def close_sync(self) -> None:
        """Close the underlying `aiohttp.ClientSession` synchronously."""
        if self._session is not None and not self._session.closed:
            loop = getattr(self._session, "_loop", None)
            if loop is not None and not loop.is_closed():
                self.run_sync(self.close())
            else:
                self._session = None

    def _url(self, endpoint: str) -> str:
        if endpoint.startswith(("http://", "https://")):
            return endpoint
        return f"{self.base_url}/{endpoint.lstrip('/')}" if endpoint else self.base_url

    async def request(self, method: str, endpoint: str, **kwargs: Any) -> Any:
        """Send an async HTTP request and return decoded JSON (via `msgspec`) or text."""
        session = await self.get_session()
        url = self._url(endpoint)
        async with self._get_semaphore():
            t_req = time.perf_counter()
            try:
                async with session.request(method, url, **kwargs) as resp:
                    raw = await resp.read()
                    if resp.status >= 400:
                        body = raw.decode("utf-8", errors="replace").strip()
                        reason = resp.reason or "HTTP Error"
                        detail = f"{reason}: {body[:300]}" if body else reason
                        raise aiohttp.ClientResponseError(
                            resp.request_info,
                            resp.history,
                            status=resp.status,
                            message=detail,
                            headers=resp.headers,
                        )
                    ctype = (resp.content_type or "").lower()
                    if "json" in ctype:
                        return msgspec.json.decode(raw)
                    return raw.decode("utf-8", errors="replace")
            finally:
                await self.polite_sleep(t_req)

    async def get(self, endpoint: str, **kwargs: Any) -> Any:
        """Perform an async HTTP GET request."""
        return await self.request("GET", endpoint, **kwargs)

    async def post(self, endpoint: str, **kwargs: Any) -> Any:
        """Perform an async HTTP POST request."""
        return await self.request("POST", endpoint, **kwargs)

    async def delete(self, endpoint: str, **kwargs: Any) -> Any:
        """Perform an async HTTP DELETE request."""
        return await self.request("DELETE", endpoint, **kwargs)

    async def stream_bytes(self, method: str, endpoint: str, **kwargs: Any) -> bytes:
        """Send an async HTTP request and read raw response bytes."""
        session = await self.get_session()
        url = self._url(endpoint)
        async with self._get_semaphore():
            t_req = time.perf_counter()
            try:
                async with session.request(method, url, **kwargs) as resp:
                    resp.raise_for_status()
                    return await resp.read()
            finally:
                await self.polite_sleep(t_req)

    async def get_json[T](self, endpoint: str, *, type_: type[T], **kwargs: Any) -> T:
        """Fetch `endpoint` via GET and decode JSON directly into `type_` using `msgspec`."""
        raw = await self.stream_bytes("GET", endpoint, **kwargs)
        return msgspec.json.decode(raw, type=type_)

    async def post_json[T](self, endpoint: str, *, type_: type[T], **kwargs: Any) -> T:
        """Fetch `endpoint` via POST and decode JSON directly into `type_` using `msgspec`."""
        raw = await self.stream_bytes("POST", endpoint, **kwargs)
        return msgspec.json.decode(raw, type=type_)

    async def stream_buffer(self, method: str, endpoint: str, **kwargs: Any) -> BytesIO:
        """Send an async HTTP request and return an in-memory `BytesIO` stream."""
        b = await self.stream_bytes(method, endpoint, **kwargs)
        return BytesIO(b)

    async def stream_chunks(
        self,
        method: str,
        endpoint: str,
        chunk_size: int = 65536,
        **kwargs: Any,
    ) -> AsyncIterator[bytes]:
        """Stream response bytes in chunks as an async generator, emitting progress when `Content-Length` is known."""
        session = await self.get_session()
        url = self._url(endpoint)
        prog = current_progress()
        async with session.request(method, url, **kwargs) as resp:
            resp.raise_for_status()
            total_bytes = resp.content_length if resp.content_length and resp.content_length > 0 else None
            received = 0
            async for chunk in resp.content.iter_chunked(chunk_size):
                received += len(chunk)
                if total_bytes is not None and prog.is_active:
                    prog.report(
                        min(1.0, received / total_bytes),
                        detail=f"Streaming {endpoint} ({received:,} / {total_bytes:,} B)",
                    )
                yield chunk

    def run_sync[T](self, coro: Coroutine[Any, Any, T]) -> T:
        """Execute a coroutine synchronously and close any session bound to the temporary event loop."""

        async def _runner() -> T:
            try:
                return await coro
            finally:
                await self.close()

        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None

        if loop and loop.is_running():
            with ThreadPoolExecutor(max_workers=1) as pool:
                future = pool.submit(lambda: asyncio.run(_runner()))
                return future.result()
        return asyncio.run(_runner())

@dataclass(frozen=True, slots=True)
class NuccoreSummary:
    """Immutable NCBI `nuccore` `esummary` record."""

    uid: str
    accession_version: str
    title: str = ""
    subname: str = ""
    subtype: str = ""
    slen: int = 0

    @property
    def base_accession(self) -> str:
        return self.accession_version.split(".")[0] if self.accession_version else ""

    @property
    def blob_text(self) -> str:
        return f"{self.title} | {self.subname}"

    @property
    def sub_pairs(self) -> list[tuple[str, str]]:
        return list(zip(self.subtype.split("|"), self.subname.split("|"), strict=False))

@dataclass(frozen=True, slots=True)
class PubMedArticleId:
    """Article identifier entry inside a PubMed `esummary` record."""

    idtype: str = ""
    value: str = ""

@dataclass(frozen=True, slots=True)
class PubMedSummary:
    """Immutable PubMed `esummary` record."""

    uid: str
    title: str = ""
    pubdate: str = ""
    doi: str = ""
    articleids: tuple[PubMedArticleId, ...] = ()

@dataclass(frozen=True, slots=True)
class DatasetAssemblyReport:
    """Immutable NCBI Datasets v2 genome assembly report summary."""

    accession: str
    organism_name: str = ""
    tax_id: int = 0
    strain: str = ""
    serovar: str = ""
    isolate: str = ""
    biosample_accession: str = ""
    assembly_level: str = ""
    paired_accession: str = ""

class NCBIClient(Client):
    """Asynchronous client for NCBI Entrez E-Utilities and NCBI Datasets v2."""

    def __init__(
        self,
        api_key: str | None = None,
        *,
        timeout: float = 40.0,
        max_concurrency: int | None = None,
        polite_delay: float = 0.20,
        recorder: Any | None = None,
    ) -> None:
        resolved_key = api_key if api_key is not None else os.environ.get("NCBI_API_KEY", "")
        conc = max_concurrency if max_concurrency is not None else (5 if resolved_key else 3)
        super().__init__(
            base_url=NCBI_EUTILS_URL,
            api_key_header="api-key" if resolved_key else None,
            api_key=resolved_key or None,
            user_agent="seroverse-ncbi-client/0.1",
            timeout=timeout,
            max_concurrency=conc,
            polite_delay=polite_delay,
            recorder=recorder,
        )

    def _eutils_params(self, extra: dict[str, str]) -> dict[str, str]:
        base = {"retmode": "json"}
        if self.api_key:
            base["api_key"] = self.api_key
        base.update(extra)
        return base

    async def esearch(self, db: str, term: str, retmax: int = 100) -> list[str]:
        """Run rate-limited `esearch.fcgi` and return matching UID strings."""
        try:
            params = self._eutils_params({"db": db, "term": term, "retmax": str(retmax)})
            env = await self.get_json("esearch.fcgi", type_=_ESearchEnvelope, params=params)
            if env.esearchresult and env.esearchresult.idlist:
                return [str(u) for u in env.esearchresult.idlist]
        except Exception:
            pass
        return []

    async def esummary(self, db: str, uids: Sequence[str], chunk_size: int = 200) -> dict[str, dict[str, Any]]:
        """Run batched `esummary.fcgi` POST requests and return `{uid: info_dict}`."""
        out: dict[str, dict[str, Any]] = {}
        chunks = [uids[i : i + chunk_size] for i in range(0, len(uids), chunk_size)]
        if not chunks:
            return out
        prog = current_progress()
        completed = 0
        total = len(chunks)

        async def _fetch_chunk(chunk: Sequence[str]) -> None:
            nonlocal completed
            try:
                data = self._eutils_params({"db": db, "id": ",".join(chunk)})
                raw = await self.stream_bytes("POST", "esummary.fcgi", data=data)
                payload = msgspec.json.decode(raw)
                if isinstance(payload, dict) and isinstance(res := payload.get("result"), dict):
                    for u in chunk:
                        if isinstance(info := res.get(u), dict):
                            out[u] = info
            except Exception:
                pass
            finally:
                completed += 1
                if prog.is_active:
                    prog.report(completed / total, detail=f"NCBI esummary ({db}): batch {completed}/{total}")

        await asyncio.gather(*[_fetch_chunk(c) for c in chunks])
        return out

    async def fetch_nuccore_summaries(
        self,
        uids: Sequence[str],
        *,
        slen_range: tuple[int, int] | None = None,
        chunk_size: int = 200,
    ) -> dict[str, NuccoreSummary]:
        """Fetch `NuccoreSummary` records for `uids`, optionally filtering by sequence length `slen_range`."""
        raw_map = await self.esummary("nuccore", uids, chunk_size=chunk_size)
        summaries: dict[str, NuccoreSummary] = {}
        for u, info in raw_map.items():
            slen = int(info.get("slen") or 0)
            if slen_range is not None and not (slen_range[0] <= slen <= slen_range[1]):
                continue
            acc_ver = str(info.get("accessionversion") or info.get("caption") or "").strip()
            if acc_ver:
                summaries[u] = NuccoreSummary(
                    uid=u,
                    accession_version=acc_ver,
                    title=str(info.get("title") or ""),
                    subname=str(info.get("subname") or ""),
                    subtype=str(info.get("subtype") or ""),
                    slen=slen,
                )
        return summaries

    async def fetch_pubmed_summaries(
        self, pmids: Sequence[int | str], chunk_size: int = 150
    ) -> dict[int, PubMedSummary]:
        """Fetch `PubMedSummary` records (including DOIs) for `pmids`."""
        uid_strs = [str(p) for p in pmids]
        raw_map = await self.esummary("pubmed", uid_strs, chunk_size=chunk_size)
        out: dict[int, PubMedSummary] = {}
        for u_s, info in raw_map.items():
            if not u_s.isdigit():
                continue
            aids: list[PubMedArticleId] = []
            doi = ""
            for aid in info.get("articleids") or ():
                if isinstance(aid, dict):
                    idt = str(aid.get("idtype") or "")
                    val = str(aid.get("value") or "").strip()
                    aids.append(PubMedArticleId(idtype=idt, value=val))
                    if idt == "doi" and val and not doi:
                        doi = val
            out[int(u_s)] = PubMedSummary(
                uid=u_s,
                title=str(info.get("title") or ""),
                pubdate=str(info.get("pubdate") or ""),
                doi=doi,
                articleids=tuple(aids),
            )
        return out

    async def elink_pubmed_to_nuccore(self, pmids: Sequence[int], batch_size: int = 40) -> dict[int, list[str]]:
        """Map PubMed IDs to linked `nuccore` UIDs via `elink.fcgi` (`pubmed_nuccore`)."""
        out: dict[int, list[str]] = {}
        batches = [pmids[i : i + batch_size] for i in range(0, len(pmids), batch_size)]
        if not batches:
            return out
        prog = current_progress()
        total = len(batches)
        for idx, batch in enumerate(batches, start=1):
            try:
                params: list[tuple[str, str]] = [
                    ("dbfrom", "pubmed"),
                    ("db", "nuccore"),
                    ("linkname", "pubmed_nuccore"),
                    ("retmode", "json"),
                ]
                if self.api_key:
                    params.append(("api_key", self.api_key))
                for p in batch:
                    params.append(("id", str(p)))
                raw = await self.stream_bytes("GET", "elink.fcgi", params=params)
                payload = msgspec.json.decode(raw)
                if not isinstance(payload, dict):
                    continue
                for lset in payload.get("linksets") or ():
                    if not isinstance(lset, dict):
                        continue
                    ids = lset.get("ids") or ()
                    if not ids:
                        continue
                    pmid = int(ids[0])
                    uids: list[str] = []
                    for ldb in lset.get("linksetdbs") or ():
                        if isinstance(ldb, dict) and ldb.get("linkname") == "pubmed_nuccore":
                            uids.extend(str(x) for x in (ldb.get("links") or ()))
                    if uids:
                        out[pmid] = uids
            except Exception:
                pass
            finally:
                if prog.is_active:
                    prog.report(idx / total, detail=f"NCBI elink: batch {idx}/{total}")
        return out

    async def efetch_fasta(self, uids: Sequence[str], db: str = "nuccore") -> bytes:
        """Download multi-FASTA bytes for `uids` via `efetch.fcgi`."""
        if not uids:
            return b""
        try:
            params = {"db": db, "id": ",".join(uids), "rettype": "fasta", "retmode": "text"}
            if self.api_key:
                params["api_key"] = self.api_key
            return await self.stream_bytes("GET", "efetch.fcgi", params=params)
        except Exception:
            return b""

    async def fetch_taxon_assemblies(self, taxon: str, page_size: int = 1000) -> list[DatasetAssemblyReport]:
        """Query NCBI Datasets v2 `genome/taxon/{taxon}/dataset_report` and return assembly reports."""
        import urllib.parse

        encoded = urllib.parse.quote(taxon)
        url = f"{NCBI_DATASETS_URL}/genome/taxon/{encoded}/dataset_report"
        try:
            raw = await self.stream_bytes("GET", url, params={"page_size": str(page_size)})
            payload = msgspec.json.decode(raw)
        except Exception:
            return []
        if not isinstance(payload, dict):
            return []
        reports: list[DatasetAssemblyReport] = []
        for rep in payload.get("reports") or ():
            if not isinstance(rep, dict):
                continue
            acc = str(rep.get("accession") or "").strip()
            if not acc:
                continue
            org = rep.get("organism") or {}
            infraspec = (org.get("infraspecific_names") or {}) if isinstance(org, dict) else {}
            asm_info = rep.get("assembly_info") or {}
            biosample = (asm_info.get("biosample") or {}) if isinstance(asm_info, dict) else {}
            strain = str(infraspec.get("strain") or biosample.get("strain") or "").strip()
            serovar = str(infraspec.get("serovar") or biosample.get("serovar") or "").strip()
            isolate = str(infraspec.get("isolate") or biosample.get("isolate") or "").strip()
            reports.append(
                DatasetAssemblyReport(
                    accession=acc,
                    organism_name=str(org.get("organism_name") or "").strip() if isinstance(org, dict) else "",
                    tax_id=int(org.get("tax_id") or 0) if isinstance(org, dict) else 0,
                    strain=strain,
                    serovar=serovar,
                    isolate=isolate,
                    biosample_accession=str(biosample.get("accession") or "").strip(),
                    assembly_level=str(asm_info.get("assembly_level") or "").strip()
                    if isinstance(asm_info, dict)
                    else "",
                    paired_accession=str(rep.get("paired_accession") or "").strip(),
                )
            )
        return reports

    async def download_assembly_fasta(self, accession: str) -> bytes:
        """Download genome FASTA bytes for `accession` via NCBI Datasets v2 archive endpoint."""
        url = f"{NCBI_DATASETS_URL}/genome/accession/{accession}/download"
        try:
            zip_bytes = await self.stream_bytes("GET", url, params={"include_annotation_type": "GENOME_FASTA"})
            with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
                fna_chunks = [
                    zf.read(zi.filename) for zi in zf.infolist() if zi.filename.endswith((".fna", ".fasta", ".fa"))
                ]
                return b"\n".join(fna_chunks)
        except Exception:
            return b""