codonyat 1.2.0__tar.gz → 1.2.1__tar.gz

{codonyat-1.2.0/codonyat.egg-info → codonyat-1.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: codonyat
-Version: 1.2.0
+Version: 1.2.1
 Summary: codon-yat — A codon-aware (Yet Another) amino acid variant typer from SAM alignments of viral NGS data.
 Author-email: Marc Noguera Julian <info@treetopunder.com>
 License-Expression: MIT
@@ -43,7 +43,7 @@ A Python package for codon-level amino acid variant calling from SAM alignments
 - Produces consistent TSV/XML diagnostics while exposing Python objects (`SamContainer`, `FullReference`, etc.) that downstream tooling can import.
 - Validates every input file before parsing to surface malformed SAM/FASTA/amplicon data early.
 - Logs per-position entropy and strand balance for easier debugging in CI or local runs.
-- Supports any protein target via `--protein` (defaults to RT).
+- Supports any protein target via `--protein` (one name, a comma-separated list, or `all`; defaults to RT).
 
 ## Installation
 
@@ -79,7 +79,7 @@ codonyat sample.sam reference.fasta amplicons.tsv --annotation annotations.tsv -
 | Flag | Default | Description |
 |------|---------|-------------|
 | `--annotation` | None | Optional TSV annotation file (record, protein, description, intervals, strand) |
-| `--protein` | `RT` | Target protein name (must match a FASTA header or TSV annotation) |
+| `--protein` | `RT` | Protein(s) to analyse: one name, a comma-separated list (e.g. `PR,RT,INT`, counted in a single pass) or `all` (every protein annotated in the reference header or TSV). Names must match the FASTA header annotations |
 | `--ratio-upper` | `3.162` | Upper bound for strand-ratio balancing |
 | `--ratio-lower` | `0.316` | Lower bound for strand-ratio balancing |
 | `--entropy-threshold` | `0.0` | Minimum Shannon entropy to mark a position as diverse |
@@ -94,6 +94,7 @@ The CLI writes two files alongside the input SAM:
   - Rows are ordered by genomic position (codon index) along the protein's intervals
   - For multi-interval proteins, this follows the biological reading order through all segments
 - `[sam-file].xml` — per-position `<Depth>`, `<FwCover>`, `<RvCover>`, `<Variants>`
+  - With several proteins, each protein's positions are wrapped in `<Protein name="…">`; single-protein XML is unchanged
 
 ### Reference FASTA format
 
@@ -152,7 +153,7 @@ result = call_variants(
     sam_path=Path("reads.sam"),
     reference_path=Path("reference.fasta"),
     amplicons_path=Path("amps.tsv"),
-    protein="RT",
+    protein="PR,RT,INT",  # or "all", "RT", or ["PR", "RT"]
     ratio_upper=3.2,
 )
 
@@ -181,7 +182,7 @@ call_variants_from_args(args)
 A lightweight runner at `codonyat-runner` exposes the same arguments as `call_variants_from_args`:
 
 ```bash
-codonyat-runner reads.sam reference.fasta amps.tsv --protein RT --ratio-upper 3.1 --csv-path results.tsv
+codonyat-runner reads.sam reference.fasta amps.tsv --protein PR,RT,INT --ratio-upper 3.1 --csv-path results.tsv
 ```
 
 ## Development

{codonyat-1.2.0 → codonyat-1.2.1}/README.md

@@ -14,7 +14,7 @@ A Python package for codon-level amino acid variant calling from SAM alignments
 - Produces consistent TSV/XML diagnostics while exposing Python objects (`SamContainer`, `FullReference`, etc.) that downstream tooling can import.
 - Validates every input file before parsing to surface malformed SAM/FASTA/amplicon data early.
 - Logs per-position entropy and strand balance for easier debugging in CI or local runs.
-- Supports any protein target via `--protein` (defaults to RT).
+- Supports any protein target via `--protein` (one name, a comma-separated list, or `all`; defaults to RT).
 
 ## Installation
 
@@ -50,7 +50,7 @@ codonyat sample.sam reference.fasta amplicons.tsv --annotation annotations.tsv -
 | Flag | Default | Description |
 |------|---------|-------------|
 | `--annotation` | None | Optional TSV annotation file (record, protein, description, intervals, strand) |
-| `--protein` | `RT` | Target protein name (must match a FASTA header or TSV annotation) |
+| `--protein` | `RT` | Protein(s) to analyse: one name, a comma-separated list (e.g. `PR,RT,INT`, counted in a single pass) or `all` (every protein annotated in the reference header or TSV). Names must match the FASTA header annotations |
 | `--ratio-upper` | `3.162` | Upper bound for strand-ratio balancing |
 | `--ratio-lower` | `0.316` | Lower bound for strand-ratio balancing |
 | `--entropy-threshold` | `0.0` | Minimum Shannon entropy to mark a position as diverse |
@@ -65,6 +65,7 @@ The CLI writes two files alongside the input SAM:
   - Rows are ordered by genomic position (codon index) along the protein's intervals
   - For multi-interval proteins, this follows the biological reading order through all segments
 - `[sam-file].xml` — per-position `<Depth>`, `<FwCover>`, `<RvCover>`, `<Variants>`
+  - With several proteins, each protein's positions are wrapped in `<Protein name="…">`; single-protein XML is unchanged
 
 ### Reference FASTA format
 
@@ -123,7 +124,7 @@ result = call_variants(
     sam_path=Path("reads.sam"),
     reference_path=Path("reference.fasta"),
     amplicons_path=Path("amps.tsv"),
-    protein="RT",
+    protein="PR,RT,INT",  # or "all", "RT", or ["PR", "RT"]
     ratio_upper=3.2,
 )
 
@@ -152,7 +153,7 @@ call_variants_from_args(args)
 A lightweight runner at `codonyat-runner` exposes the same arguments as `call_variants_from_args`:
 
 ```bash
-codonyat-runner reads.sam reference.fasta amps.tsv --protein RT --ratio-upper 3.1 --csv-path results.tsv
+codonyat-runner reads.sam reference.fasta amps.tsv --protein PR,RT,INT --ratio-upper 3.1 --csv-path results.tsv
 ```
 
 ## Development

{codonyat-1.2.0 → codonyat-1.2.1}/aa_caller/app.py

@@ -13,7 +13,12 @@ from .constants import (  # noqa: F401
 from .container import SamContainer  # noqa: F401
 from .genetic_code import GENETIC_CODE, codon_to_aminoacid  # noqa: F401
 from .models import Amplicon, Protein, Qual, Variant  # noqa: F401
-from .reference import FullReference  # noqa: F401
+from .reference import (  # noqa: F401
+    FullReference,
+    parse_protein_spec,
+    resolve_protein_targets,
+    resolve_proteins,
+)
 from .sam import SamEntry  # noqa: F401
 from .validators import (  # noqa: F401
     parse_amplicons,

{codonyat-1.2.0 → codonyat-1.2.1}/aa_caller/cli.py

@@ -11,7 +11,7 @@ from .constants import (
     DEFAULT_RATIO_UPPER,
 )
 from .container import SamContainer
-from .reference import FullReference
+from .reference import FullReference, parse_protein_spec, resolve_protein_targets
 from .validators import (
     parse_amplicons,
     validate_amplicon_file,
@@ -58,7 +58,11 @@ def main() -> None:
         "--protein",
         type=str,
         default=DEFAULT_PROTEIN,
-        help="Target protein name to analyze (default: RT).",
+        help=(
+            "Protein(s) to analyze: one name (default: RT), a comma-separated list "
+            "(e.g. PR,RT,INT; all counted in one pass over the SAM) or 'all' for "
+            "every protein annotated in the reference header."
+        ),
     )
     args = parser.parse_args()
 
@@ -73,7 +77,13 @@ def main() -> None:
 
     logger.info("Validating input formats before parsing")
     validate_sam_file(args.sam_file)
-    validate_reference_file(args.reference_file, protein_name=args.protein, annotation_path=args.annotation)
+    try:
+        parse_protein_spec(args.protein)
+    except ValueError as exc:
+        parser.error(str(exc))
+    validate_reference_file(
+        args.reference_file, protein_name=args.protein, annotation_path=args.annotation
+    )
     validate_amplicon_file(args.amplicons_file)
 
     logger.info("Reading amplicon configuration")
@@ -81,6 +91,11 @@ def main() -> None:
 
     logger.info("Loading reference")
     reference = FullReference(args.reference_file, annotation_path=args.annotation)
+    try:
+        proteins = resolve_protein_targets(args.protein, reference)
+    except ValueError as exc:
+        parser.error(str(exc))
+    protein_names = [protein.name for protein in proteins]
 
     logger.info("Parsing SAM entries")
     container = SamContainer(
@@ -91,12 +106,17 @@ def main() -> None:
     )
     container.load()
 
-    logger.info("Calculating %s variant frequencies", args.protein)
-    container.calculate_variant_frequencies(args.protein, reference)
+    logger.info("Calculating %s variant frequencies", ",".join(protein_names))
+    container.calculate_for_proteins(proteins, reference)
 
     csv_path = args.sam_file.with_suffix(args.sam_file.suffix + ".tsv")
     xml_path = args.sam_file.with_suffix(args.sam_file.suffix + ".xml")
     logger.info("Writing TSV results")
-    container.write_csv(str(args.sam_file), reference, csv_path, protein_name=args.protein)
-    logger.info("Writing XML dump")
-    container.write_xml(str(args.sam_file), reference, xml_path)
+    if len(proteins) == 1:
+        container.write_csv(str(args.sam_file), reference, csv_path, protein_name=protein_names[0])
+        logger.info("Writing XML dump")
+        container.write_xml(str(args.sam_file), reference, xml_path)
+    else:
+        container.write_csv_proteins(str(args.sam_file), reference, csv_path, protein_names)
+        logger.info("Writing XML dump")
+        container.write_xml_proteins(str(args.sam_file), reference, xml_path, protein_names)

codonyat-1.2.1/aa_caller/container.py

@@ -0,0 +1,392 @@
+from __future__ import annotations
+
+import csv
+import logging
+import math
+import xml.etree.ElementTree as ET
+from collections.abc import Sequence
+from pathlib import Path
+
+from .genetic_code import codon_to_aminoacid
+from .models import Protein, Variant
+from .reference import FullReference
+from .sam import SamEntry
+
+logger = logging.getLogger(__name__)
+
+
+class SamContainer:
+    """Aggregates SAM entries and computes RT variant statistics for an amplicon set."""
+
+    def __init__(
+        self,
+        sam_path: Path,
+        ratio_upper: float,
+        ratio_lower: float,
+        entropy_threshold: float,
+    ) -> None:
+        """Initialize the container with the SAM path, ratio window, and entropy threshold."""
+        self.sam_path = sam_path
+        self.reads: list[SamEntry] = []
+        self.variants: dict[int, dict] = {}
+        self.protein_variants: dict[str, dict[int, dict]] = {}
+        self._protein_result_list: list[tuple[str, dict[int, dict]]] = []
+        self.ratio_upper = ratio_upper
+        self.ratio_lower = ratio_lower
+        self.entropy_threshold = entropy_threshold
+
+    def load(self) -> None:
+        """Load every mapped read from the SAM file so they can be aggregated."""
+        with open(self.sam_path) as fh:
+            for line in fh:
+                if line.startswith("@"):
+                    continue
+                entry = SamEntry(line)
+                if entry.is_mapped():
+                    self.reads.append(entry)
+
+    def calculate_variant_frequencies(
+        self, protein_name: str, reference: FullReference, record_id: str | None = None
+    ) -> None:
+        """Walk every position of the requested protein and build variant histories.
+
+        Equivalent to ``calculate_variant_frequencies_multi([protein_name], reference)``
+        when ``record_id`` is omitted. Supports multi-interval proteins
+        (frameshift, splicing, circular).
+        """
+        protein = reference.get_protein(protein_name, record_id)
+        if protein is None:
+            raise ValueError(f"Protein {protein_name} not found in reference")
+        self.calculate_for_proteins([protein], reference)
+
+    def calculate_variant_frequencies_multi(
+        self, protein_names: Sequence[str], reference: FullReference
+    ) -> None:
+        """Count codons for several proteins in a single pass over loaded reads.
+
+        Results are stored per protein in ``self.protein_variants``;
+        ``self.variants`` points at the first requested protein (the only one in
+        single-protein runs), as in 1.1.0.
+        """
+        if not protein_names:
+            raise ValueError("At least one protein is required")
+        proteins: list[Protein] = []
+        for name in protein_names:
+            protein = reference.get_protein(name)
+            if protein is None:
+                raise ValueError(f"Protein {name} not found in reference")
+            proteins.append(protein)
+        self.calculate_for_proteins(proteins, reference)
+
+    def calculate_for_proteins(self, proteins: Sequence[Protein], reference: FullReference) -> None:
+        """Count every requested protein in one pass over ``self.reads``.
+
+        Supports multi-record references (reads are matched by RNAME) and
+        multi-interval proteins.
+        """
+        if not proteins:
+            raise ValueError("At least one protein is required")
+
+        jobs: list[tuple[Protein, dict[int, dict]]] = []
+        for protein in proteins:
+            entries: dict[int, dict] = {}
+            for pos, record_id in self._get_codon_positions(protein, reference):
+                entries[pos] = {
+                    "variants": {},
+                    "fw_depth": 0,
+                    "rv_depth": 0,
+                    "depth": 0,
+                    "amplicons": {},
+                    "record_id": record_id,
+                }
+            jobs.append((protein, entries))
+
+        for read in self.reads:
+            for protein, entries in jobs:
+                if read.reference != protein.record_id:
+                    continue
+                for pos, entry in entries.items():
+                    self._accumulate_read_at_position(read, pos, entry)
+
+        self.protein_variants = {}
+        self._protein_result_list = []
+        for protein, entries in jobs:
+            self.variants = entries
+            for pos in self.variants:
+                self.calculate_ratios_nucleotide(pos)
+                self.calculate_ratios_aminoacid(pos)
+                self.calculate_shannon_entropy_single_position(pos)
+                self.log_position_stats(pos)
+            self.protein_variants[protein.name] = self.variants
+            self._protein_result_list.append((protein.name, self.variants))
+        self.variants = jobs[0][1]
+
+    def _accumulate_read_at_position(self, read: SamEntry, pos: int, entry: dict) -> None:
+        """Add one mapped read's codon to a protein position, if it covers it."""
+        if not read.covers(pos):
+            return
+        codon = read.codon_at(pos)
+        if not codon or len(codon) < 3:
+            return
+        depth_increment = read.occurences
+        variants: dict[str, Variant] = entry["variants"]
+        if read.orientation == "F":
+            entry["fw_depth"] += depth_increment
+        else:
+            entry["rv_depth"] += depth_increment
+        entry["depth"] += depth_increment
+        variant = variants.setdefault(codon, Variant(codon=codon))
+        variant.count += depth_increment
+        if read.orientation == "F":
+            variant.fw_reads += depth_increment
+        else:
+            variant.rv_reads += depth_increment
+        amplicon_label = read.amplicon or "Amp_NONE"
+        amp_variant = variant.amplicons.setdefault(amplicon_label, {"fw": 0, "rv": 0})
+        if read.orientation == "F":
+            amp_variant["fw"] += depth_increment
+        else:
+            amp_variant["rv"] += depth_increment
+        amp_position = entry["amplicons"].setdefault(
+            amplicon_label, {"fw_depth": 0, "rv_depth": 0, "depth": 0}
+        )
+        if read.orientation == "F":
+            amp_position["fw_depth"] += depth_increment
+        else:
+            amp_position["rv_depth"] += depth_increment
+        amp_position["depth"] += depth_increment
+
+    def _get_codon_positions(self, protein, reference: FullReference) -> list[tuple]:
+        """Generate list of (position, record_id) tuples for all codons in the protein.
+
+        Handles multi-interval proteins including frameshifts.
+        """
+        positions = []
+        current_offset = 0
+
+        for interval in protein.intervals:
+            interval_length = interval.end - interval.start + 1
+            pos = interval.start
+
+            while pos <= interval.end:
+                remaining_in_interval = interval.end - pos + 1
+                if remaining_in_interval >= 3:
+                    positions.append((pos, protein.record_id))
+                    pos += 3
+                else:
+                    break
+
+            current_offset += interval_length
+
+        return positions
+
+    def calculate_ratios_nucleotide(self, position: int) -> None:
+        """Determine strand-specific ratios per variant and flag the balanced reads."""
+        entry = self.variants.get(position)
+        if not entry or entry["fw_depth"] == 0 or entry["rv_depth"] == 0:
+            return
+        for variant in entry["variants"].values():
+            # compute the per-strand frequency for each codon then derive the ratio
+            variant.fw_freq = variant.fw_reads / entry["fw_depth"] if entry["fw_depth"] else 0
+            variant.rv_freq = variant.rv_reads / entry["rv_depth"] if entry["rv_depth"] else 0
+            variant.ratio = (variant.fw_freq / variant.rv_freq) if variant.rv_freq else 0
+            # only mark variants within the balanced ratio window
+            if self.ratio_lower <= variant.ratio <= self.ratio_upper:
+                variant.balanced_fw += variant.fw_reads
+                variant.balanced_rv += variant.rv_reads
+
+    def has_balanced_variants(self, position: int) -> bool:
+        """Return True when at least one variant falls within the configured strand-ratio window."""
+        entry = self.variants.get(position)
+        if not entry:
+            return False
+        return any(self.ratio_lower <= variant.ratio <= self.ratio_upper for variant in entry["variants"].values())
+
+    def calculate_ratios_aminoacid(self, position: int) -> None:
+        """Aggregate amino-acid frequencies across codon variants and capture balance metadata."""
+        entry = self.variants.get(position)
+        if not entry:
+            return
+        aa_stats: dict[str, dict[str, float]] = {}
+        for variant in entry["variants"].values():
+            aa = codon_to_aminoacid(variant.codon)
+            stats = aa_stats.setdefault(aa, {"count": 0.0, "fw": 0.0, "rv": 0.0})
+            stats["count"] += variant.count
+            stats["fw"] += variant.fw_reads
+            stats["rv"] += variant.rv_reads
+        for aa, stats in aa_stats.items():
+            fw = stats["fw"]
+            rv = stats["rv"]
+            ratio = (fw / rv) if rv else 0
+            # flag balanced amino-acid counts using the same ratio window as nucleotides
+            stats["ratio"] = ratio
+            stats["balanced"] = (fw + rv) if self.ratio_lower <= ratio <= self.ratio_upper else 0
+        entry["aminoacid_stats"] = aa_stats
+
+    def calculate_shannon_entropy_single_position(self, position: int) -> float:
+        """Compute the Shannon entropy for the observed codon mixture at a position."""
+        entry = self.variants.get(position)
+        if not entry:
+            return 0.0
+        depth = entry.get("depth", 0)
+        if depth == 0:
+            entry["shannon_entropy"] = 0.0
+            return 0.0
+        entropy = 0.0
+        for variant in entry["variants"].values():
+            if variant.count == 0:
+                continue
+            p = variant.count / depth
+            # standard Shannon formula: sum of -p log(p) for each variant
+            entropy -= p * math.log(p)
+        entry["shannon_entropy"] = entropy
+        entry["entropy_pass"] = entropy >= self.entropy_threshold
+        return entropy
+
+    def log_position_stats(self, position: int) -> None:
+        """Log strand balance and entropy so downstream debugging is easier."""
+        entry = self.variants.get(position)
+        if not entry:
+            return
+        balanced = [v.codon for v in entry["variants"].values() if self.ratio_lower <= v.ratio <= self.ratio_upper]
+        entropy = entry.get("shannon_entropy", 0.0)
+        logger.debug(
+            "Position %s -- depth=%s, fw=%s, rv=%s, entropy=%.3f, balanced=%s",
+            position,
+            entry["depth"],
+            entry["fw_depth"],
+            entry["rv_depth"],
+            entropy,
+            balanced,
+        )
+
+    def _entries_for(self, protein_name: str) -> dict[int, dict]:
+        """Per-position results for a protein (falls back to ``self.variants``)."""
+        return self.protein_variants.get(protein_name, self.variants)
+
+    def _csv_fieldnames(self, reference: FullReference) -> list[str]:
+        has_multiple_records = len(reference.record_ids) > 1
+        fieldnames = ["FILE", "REFERENCE"]
+        if has_multiple_records:
+            fieldnames.append("RECORD")
+        fieldnames.extend(
+            ["PROTEIN", "VARIANT", "POSITION", "FREQ", "FWCOV", "RVCOV", "TOTALCOV", "RATIO"]
+        )
+        return fieldnames
+
+    def _write_csv_rows(
+        self,
+        writer: csv.DictWriter,
+        sample_name: str,
+        reference: FullReference,
+        protein_name: str,
+        entries: dict[int, dict] | None = None,
+    ) -> None:
+        has_multiple_records = len(reference.record_ids) > 1
+        if entries is None:
+            entries = self._entries_for(protein_name)
+        for pos, entry in sorted(entries.items()):
+            depth = entry["depth"]
+            ratio = entry["fw_depth"] / entry["rv_depth"] if entry["rv_depth"] else 0
+            record_id = entry.get("record_id", reference.id)
+            for variant in entry["variants"].values():
+                row = {
+                    "FILE": sample_name,
+                    "REFERENCE": record_id,
+                    "PROTEIN": protein_name,
+                    "VARIANT": variant.codon,
+                    "POSITION": pos,
+                    "FREQ": round((variant.count / depth * 100) if depth else 0, 3),
+                    "FWCOV": entry["fw_depth"],
+                    "RVCOV": entry["rv_depth"],
+                    "TOTALCOV": depth,
+                    "RATIO": round(ratio, 3),
+                }
+                if has_multiple_records:
+                    row["RECORD"] = record_id
+                writer.writerow(row)
+
+    def write_csv(
+        self, sample_name: str, reference: FullReference, output_path: Path, *, protein_name: str = "RT"
+    ) -> None:
+        """Emit the variant summary as a TSV file that mirrors the Perl outputs.
+
+        For multi-record references, adds a RECORD column after REFERENCE.
+        Single-record outputs remain byte-identical to 1.1.0.
+        """
+        self.write_csv_proteins(sample_name, reference, output_path, [protein_name])
+
+    def write_csv_proteins(
+        self,
+        sample_name: str,
+        reference: FullReference,
+        output_path: Path,
+        protein_names: Sequence[str],
+    ) -> None:
+        """Emit one TSV with a single header and the rows of each protein, in the given order."""
+        fieldnames = self._csv_fieldnames(reference)
+        pairs = self._output_pairs(protein_names)
+        with open(output_path, "w", newline="") as csvfile:
+            writer = csv.DictWriter(csvfile, fieldnames=fieldnames, delimiter="\t")
+            writer.writeheader()
+            for protein_name, entries in pairs:
+                self._write_csv_rows(writer, sample_name, reference, protein_name, entries)
+
+    def _output_pairs(self, protein_names: Sequence[str]) -> list[tuple[str, dict[int, dict]]]:
+        """Map requested names to counted entries, preserving duplicate-name order."""
+        calculated = [name for name, _entries in self._protein_result_list]
+        if self._protein_result_list and calculated == list(protein_names):
+            return list(self._protein_result_list)
+        used: set[int] = set()
+        pairs: list[tuple[str, dict[int, dict]]] = []
+        for name in protein_names:
+            found = False
+            for index, (result_name, entries) in enumerate(self._protein_result_list):
+                if result_name == name and index not in used:
+                    used.add(index)
+                    pairs.append((result_name, entries))
+                    found = True
+                    break
+            if not found:
+                pairs.append((name, self._entries_for(name)))
+        return pairs
+
+    @staticmethod
+    def _add_positions(parent: ET.Element, entries: dict[int, dict]) -> None:
+        for pos, entry in sorted(entries.items()):
+            position_elem = ET.SubElement(parent, "Position", index=str(pos))
+            ET.SubElement(position_elem, "Depth").text = str(entry["depth"])
+            ET.SubElement(position_elem, "FwCover").text = str(entry["fw_depth"])
+            ET.SubElement(position_elem, "RvCover").text = str(entry["rv_depth"])
+            variants_elem = ET.SubElement(position_elem, "Variants")
+            for variant in entry["variants"].values():
+                # describe each variant using XML attributes for diagnostics
+                var_elem = ET.SubElement(variants_elem, "Variant", codon=variant.codon)
+                var_elem.set("count", str(variant.count))
+                var_elem.set("fw_reads", str(variant.fw_reads))
+                var_elem.set("rv_reads", str(variant.rv_reads))
+
+    def write_xml(self, sample_name: str, reference: FullReference, output_path: Path) -> None:
+        """Dump the per-position state of the (first) protein into XML for diagnostics."""
+        root = ET.Element("SamContainer", sample=sample_name, reference=reference.id)
+        self._add_positions(root, self.variants)
+        tree = ET.ElementTree(root)
+        tree.write(output_path, encoding="utf-8", xml_declaration=True)
+
+    def write_xml_proteins(
+        self,
+        sample_name: str,
+        reference: FullReference,
+        output_path: Path,
+        protein_names: Sequence[str],
+    ) -> None:
+        """Multi-protein XML: one ``<Protein name=...>`` element per protein wrapping its positions.
+
+        For a single protein use :meth:`write_xml`, whose layout is unchanged.
+        """
+        root = ET.Element("SamContainer", sample=sample_name, reference=reference.id)
+        for protein_name, entries in self._output_pairs(protein_names):
+            protein_elem = ET.SubElement(root, "Protein", name=protein_name)
+            self._add_positions(protein_elem, entries)
+        tree = ET.ElementTree(root)
+        tree.write(output_path, encoding="utf-8", xml_declaration=True)

{codonyat-1.2.0 → codonyat-1.2.1}/aa_caller/reference.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import csv
+from collections.abc import Sequence
 from pathlib import Path
 
 from Bio import SeqIO
@@ -139,3 +140,77 @@ class FullReference:
 
         seq = self.sequences.get(record_id, "")
         return seq[position - 1 : position - 1 + length]
+
+    def proteins_in_order(self) -> list[Protein]:
+        """Return annotated proteins in FASTA-header / TSV order."""
+        return list(self.proteins.values())
+
+
+def parse_protein_spec(spec: str | Sequence[str]) -> list[str]:
+    """Split a protein selection (``"RT"``, ``"PR,RT,INT"``, ``"all"`` or a list) into names.
+
+    Order is kept and duplicates are dropped. ``"all"`` is returned as ``["all"]``
+    and expanded by :func:`resolve_proteins` once the reference is loaded.
+    """
+    if isinstance(spec, str):
+        items = [part.strip() for part in spec.split(",")]
+    else:
+        items = [str(part).strip() for part in spec]
+    names: list[str] = []
+    for item in items:
+        if item and item not in names:
+            names.append(item)
+    if not names:
+        raise ValueError("No protein selected")
+    if "all" in (name.lower() for name in names) and len(names) > 1:
+        raise ValueError("'all' cannot be combined with other protein names")
+    return names
+
+
+def resolve_protein_targets(spec: str | Sequence[str], reference: FullReference) -> list[Protein]:
+    """Resolve a protein selection to annotated Protein objects.
+
+    ``"all"`` selects every annotated protein in FASTA-header / TSV order.
+    A comma-separated list keeps the requested order. On multi-record
+    references, each requested name includes every matching record.
+    """
+    names = parse_protein_spec(spec)
+    if len(names) == 1 and names[0].lower() == "all":
+        proteins = reference.proteins_in_order()
+        if not proteins:
+            raise ValueError(f"Reference {reference.path} has no annotated proteins")
+        return proteins
+    selected: list[Protein] = []
+    missing: list[str] = []
+    for name in names:
+        matches = [protein for key, protein in reference.proteins.items() if key[1] == name]
+        if not matches:
+            missing.append(name)
+        else:
+            selected.extend(matches)
+    if missing:
+        raise ValueError(f"Protein(s) not found in reference: {', '.join(missing)}")
+    return selected
+
+
+def resolve_proteins(spec: str | Sequence[str], reference: FullReference) -> list[str]:
+    """Resolve a protein selection to protein names.
+
+    ``"all"`` selects every annotated protein in FASTA-header / TSV order.
+    Names are unique; duplicate names across records are collapsed in
+    first-seen order. Use :func:`resolve_protein_targets` when record
+    identity must be preserved.
+    """
+    names = parse_protein_spec(spec)
+    if len(names) == 1 and names[0].lower() == "all":
+        if not reference.proteins:
+            raise ValueError(f"Reference {reference.path} has no annotated proteins")
+        seen: list[str] = []
+        for _record_id, protein_name in reference.proteins:
+            if protein_name not in seen:
+                seen.append(protein_name)
+        return seen
+    missing = [name for name in names if reference.get_protein(name) is None]
+    if missing:
+        raise ValueError(f"Protein(s) not found in reference: {', '.join(missing)}")
+    return names