123 lines
6.2 KiB
Python
123 lines
6.2 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._-]*$")
|
|
|
|
|
|
def extract_archive(archive_path: Path, destination: Path) -> None:
|
|
destination_root = destination.resolve()
|
|
with tarfile.open(archive_path, "r:*") as archive:
|
|
for member in archive.getmembers():
|
|
target = (destination / member.name).resolve()
|
|
if target != destination_root and destination_root not in target.parents:
|
|
raise ValueError(f"unsafe archive member: {member.name}")
|
|
archive.extractall(destination)
|
|
|
|
|
|
def discover(root: Path, suffix: str) -> dict[str, Path]:
|
|
found: dict[str, Path] = {}
|
|
for path in root.rglob(f"*{suffix}"):
|
|
accession = path.name[:-len(suffix)]
|
|
if not SAFE_ACCESSION.fullmatch(accession):
|
|
raise ValueError(f"unsafe accession derived from {path.name}")
|
|
if accession in found:
|
|
raise ValueError(f"duplicate {suffix} file for {accession}")
|
|
found[accession] = path
|
|
return found
|
|
|
|
|
|
def run_checked(command: list[str], accession: str, label: str) -> None:
|
|
completed = subprocess.run(command, text=True, capture_output=True)
|
|
if completed.returncode != 0:
|
|
raise RuntimeError(f"{accession} {label} failed:\n{completed.stderr or completed.stdout}")
|
|
|
|
|
|
def run_one(accession: str, proteins: Path, hmm_hits: Path, homology_hits: Path, rules: Path, candidates_dir: Path, members_dir: Path, merge_script: Path, classify_script: Path) -> dict[str, str | int]:
|
|
candidates = candidates_dir / f"{accession}.candidates.tsv"
|
|
members = members_dir / f"{accession}.final_members.tsv"
|
|
run_checked([
|
|
sys.executable, str(merge_script), "--hmm", str(hmm_hits), "--homology", str(homology_hits),
|
|
"--accession", accession, "--output", str(candidates),
|
|
], accession, "candidate merge")
|
|
run_checked([
|
|
sys.executable, str(classify_script), "--candidates", str(candidates), "--hmm-hits", str(hmm_hits),
|
|
"--homology-hits", str(homology_hits), "--proteins", str(proteins), "--rules", str(rules), "--output", str(members),
|
|
], accession, "member classification")
|
|
return {"accession": accession, "candidates": candidates.name, "final_members": members.name, "members_bytes": members.stat().st_size}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Merge and classify family candidates for every accession in the input bundles")
|
|
parser.add_argument("--primary-bundle", required=True, type=Path)
|
|
parser.add_argument("--hmm-bundle", required=True, type=Path)
|
|
parser.add_argument("--homology-bundle", required=True, type=Path)
|
|
parser.add_argument("--rules", 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")
|
|
archive_output = Path(f"{args.output_prefix}.members.tar.gz")
|
|
manifest_output = Path(f"{args.output_prefix}.members_manifest.tsv")
|
|
archive_output.parent.mkdir(parents=True, exist_ok=True)
|
|
merge_script = Path(__file__).with_name("merge_candidates.py").resolve()
|
|
classify_script = Path(__file__).with_name("classify_members.py").resolve()
|
|
with tempfile.TemporaryDirectory(prefix="members_batch_") as temp_name:
|
|
temp = Path(temp_name)
|
|
primary_root, hmm_root, homology_root = temp / "primary_input", temp / "hmm_input", temp / "homology_input"
|
|
candidates_dir, members_dir = temp / "candidates", temp / "members"
|
|
for directory in (primary_root, hmm_root, homology_root, candidates_dir, members_dir):
|
|
directory.mkdir()
|
|
extract_archive(args.primary_bundle.resolve(), primary_root)
|
|
extract_archive(args.hmm_bundle.resolve(), hmm_root)
|
|
extract_archive(args.homology_bundle.resolve(), homology_root)
|
|
proteins = discover(primary_root, ".primary.pep.fa")
|
|
hmm_hits = discover(hmm_root, ".hits.tsv")
|
|
homology_hits = discover(homology_root, ".hits.tsv")
|
|
accession_sets = {"primary": set(proteins), "hmm": set(hmm_hits), "homology": set(homology_hits)}
|
|
expected = accession_sets["primary"]
|
|
for label, observed in accession_sets.items():
|
|
if observed != expected:
|
|
missing = sorted(expected - observed)
|
|
extra = sorted(observed - expected)
|
|
raise ValueError(f"{label} bundle accession mismatch; missing={missing}, extra={extra}")
|
|
results: list[dict[str, str | int]] = []
|
|
with ThreadPoolExecutor(max_workers=min(args.jobs, len(expected))) as executor:
|
|
futures = {
|
|
executor.submit(
|
|
run_one, accession, proteins[accession], hmm_hits[accession], homology_hits[accession], args.rules.resolve(),
|
|
candidates_dir, members_dir, merge_script, classify_script,
|
|
): accession
|
|
for accession in expected
|
|
}
|
|
for future in as_completed(futures):
|
|
result = future.result()
|
|
results.append(result)
|
|
print(f"completed member classification: {result['accession']}", flush=True)
|
|
fields = ["accession", "candidates", "final_members", "members_bytes"]
|
|
with manifest_output.open("w", encoding="utf-8", newline="") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n")
|
|
writer.writeheader()
|
|
writer.writerows(sorted(results, key=lambda row: str(row["accession"])))
|
|
with tarfile.open(archive_output, "w:gz") as archive:
|
|
for file_path in sorted(candidates_dir.iterdir()):
|
|
archive.add(file_path, arcname=f"candidates/{file_path.name}")
|
|
for file_path in sorted(members_dir.iterdir()):
|
|
archive.add(file_path, arcname=f"members/{file_path.name}")
|
|
print(f"processed {len(expected)} samples")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|