Source code for arda.adapter

"""Library-facing API.

A stable, import-friendly surface for embedding arda in other Python tools.
The heavy lifting lands in Phase 2 (``arda.annotate``); this module keeps the
public signature stable.
"""

from __future__ import annotations

from typing import Iterable, Literal

SeqType = Literal["nt", "aa"]




def annotate_sequences(
    sequences: Iterable[str] | Iterable[tuple[str, str]],
    seqtype: SeqType = "nt",
    organism: str = "human",
    map_d: bool = True,
    *,
    threads: int = 0,
    map_dd: bool = True,
    major_only: bool = True,
):
    """Annotate FR/CDR regions for a batch of sequences.

    Args:
        sequences: Either raw sequence strings or ``(id, sequence)`` pairs.
        seqtype: ``"nt"`` for nucleotide input, ``"aa"`` for amino acid.
        organism: One of the supported organisms (human, mouse, rat, rabbit,
            rhesus_monkey).
        map_d: Map D segments (``d_call``/``d2_call``/``np*``) for VDJ-locus hits;
            ``False`` skips D mapping. Supports nucleotide and amino-acid input.
        map_dd: Search a tandem second D (default True); False retains single-D mapping.
        major_only: Search one major allele per gene (default); False searches all alleles.
        threads: CPU budget; 0 uses all CPUs available to the process/job. An already
            initialized Polars pool exceeding this budget raises ``ValueError``.

    Returns:
        A list of AIRR-style annotation record dicts, one per input sequence.
        Unmapped sequences retain their ID and sequence with empty annotation fields.
    """
    from .annotate.mapper import annotate_records

    pairs: list[tuple[str, str]] = []
    for i, item in enumerate(sequences):
        if isinstance(item, str):
            pairs.append((f"seq{i}", item))
        else:
            pairs.append((str(item[0]), str(item[1])))
    return annotate_records(pairs, organism=organism, seqtype=seqtype, map_d=map_d, map_dd=map_dd, threads=threads, major_only=major_only)