131 lines
5.4 KiB
Python
131 lines
5.4 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import tarfile
|
|
import tempfile
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from pathlib import Path
|
|
|
|
|
|
SAFE_ACCESSION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
|
|
REQUIRED_COLUMNS = ("accession", "gff3", "pep", "cds")
|
|
|
|
|
|
def resolve_input(value: str, samples_path: Path) -> Path:
|
|
path = Path(value)
|
|
candidates = [path] if path.is_absolute() else [samples_path.parent / path, Path.cwd() / path]
|
|
for candidate in candidates:
|
|
if candidate.is_file():
|
|
return candidate.resolve()
|
|
raise FileNotFoundError(f"input file does not exist: {value}")
|
|
|
|
|
|
def load_samples(samples_path: Path) -> list[dict[str, str | Path]]:
|
|
with samples_path.open(encoding="utf-8-sig", newline="") as handle:
|
|
reader = csv.DictReader(handle, delimiter="\t")
|
|
missing = [column for column in REQUIRED_COLUMNS if column not in (reader.fieldnames or [])]
|
|
if missing:
|
|
raise ValueError(f"samples.tsv is missing columns: {', '.join(missing)}")
|
|
rows = list(reader)
|
|
if not rows:
|
|
raise ValueError("samples.tsv contains no samples")
|
|
seen: set[str] = set()
|
|
normalized: list[dict[str, str | Path]] = []
|
|
for row in rows:
|
|
accession = (row.get("accession") or "").strip()
|
|
if not SAFE_ACCESSION.fullmatch(accession):
|
|
raise ValueError(f"unsafe accession token: {accession!r}")
|
|
if accession in seen:
|
|
raise ValueError(f"duplicate accession: {accession}")
|
|
seen.add(accession)
|
|
normalized.append({
|
|
"accession": accession,
|
|
"gff3": resolve_input((row.get("gff3") or "").strip(), samples_path),
|
|
"pep": resolve_input((row.get("pep") or "").strip(), samples_path),
|
|
"cds": resolve_input((row.get("cds") or "").strip(), samples_path),
|
|
})
|
|
return normalized
|
|
|
|
|
|
def run_one(row: dict[str, str | Path], script: Path, output_dir: Path) -> dict[str, str | int]:
|
|
accession = str(row["accession"])
|
|
prefix = output_dir / accession
|
|
command = [
|
|
sys.executable, str(script),
|
|
"--accession", accession,
|
|
"--gff3", str(row["gff3"]),
|
|
"--pep", str(row["pep"]),
|
|
"--cds", str(row["cds"]),
|
|
"--out-prefix", str(prefix),
|
|
]
|
|
completed = subprocess.run(command, text=True, capture_output=True)
|
|
if completed.returncode != 0:
|
|
raise RuntimeError(f"{accession} failed:\n{completed.stderr or completed.stdout}")
|
|
outputs = [
|
|
Path(f"{prefix}.primary.pep.fa"),
|
|
Path(f"{prefix}.primary.cds.fa"),
|
|
Path(f"{prefix}.id_mapping.tsv"),
|
|
]
|
|
for output in outputs:
|
|
if not output.is_file():
|
|
raise RuntimeError(f"{accession} did not create {output.name}")
|
|
return {
|
|
"accession": accession,
|
|
"primary_pep": outputs[0].name,
|
|
"primary_cds": outputs[1].name,
|
|
"id_mapping": outputs[2].name,
|
|
"primary_pep_bytes": outputs[0].stat().st_size,
|
|
"primary_cds_bytes": outputs[1].stat().st_size,
|
|
}
|
|
|
|
|
|
def write_manifest(rows: list[dict[str, str | int]], path: Path) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
fields = ["accession", "primary_pep", "primary_cds", "id_mapping", "primary_pep_bytes", "primary_cds_bytes"]
|
|
with path.open("w", encoding="utf-8", newline="") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n")
|
|
writer.writeheader()
|
|
writer.writerows(sorted(rows, key=lambda row: str(row["accession"])))
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Extract primary isoforms for every row in samples.tsv")
|
|
parser.add_argument("--samples", required=True, type=Path)
|
|
parser.add_argument("--jobs", type=int, default=4, help="maximum concurrent samples; does not limit sample count")
|
|
parser.add_argument("--output-prefix", required=True, type=Path)
|
|
args = parser.parse_args()
|
|
if args.jobs < 1:
|
|
parser.error("--jobs must be at least 1")
|
|
samples_path = args.samples.resolve()
|
|
rows = load_samples(samples_path)
|
|
script = Path(__file__).with_name("select_primary_isoform.py").resolve()
|
|
archive_path = Path(f"{args.output_prefix}.primary.tar.gz")
|
|
manifest_path = Path(f"{args.output_prefix}.primary_manifest.tsv")
|
|
archive_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(prefix="primary_batch_") as temp_name:
|
|
output_dir = Path(temp_name) / "primary"
|
|
output_dir.mkdir()
|
|
results: list[dict[str, str | int]] = []
|
|
with ThreadPoolExecutor(max_workers=min(args.jobs, len(rows))) as executor:
|
|
futures = {executor.submit(run_one, row, script, output_dir): str(row["accession"]) for row in rows}
|
|
for future in as_completed(futures):
|
|
result = future.result()
|
|
results.append(result)
|
|
print(f"completed primary isoforms: {result['accession']}", flush=True)
|
|
write_manifest(results, manifest_path)
|
|
with tarfile.open(archive_path, "w:gz") as archive:
|
|
for file_path in sorted(output_dir.iterdir()):
|
|
archive.add(file_path, arcname=f"primary/{file_path.name}")
|
|
print(f"processed {len(rows)} samples")
|
|
print(f"archive: {archive_path}")
|
|
print(f"manifest: {manifest_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|