make_samples_tsv.py

#!/usr/bin/env python3
import argparse
import re
import sys
from pathlib import Path
import pandas as pd

PATTERNS = [
    # 常见：xxx_R1.fq.gz / xxx_R2.fq.gz
    re.compile(r"^(?P<prefix>.+?)_(?P<read>R[12])\.(?:f(ast)?q)(?:\.gz)?$", re.IGNORECASE),
    # 备选：xxx.1.fq.gz / xxx.2.fq.gz
    re.compile(r"^(?P<prefix>.+?)\.(?P<num>[12])\.(?:f(ast)?q)(?:\.gz)?$", re.IGNORECASE),
    # 备选：xxx_1.fastq.gz / xxx_2.fastq.gz
    re.compile(r"^(?P<prefix>.+?)_(?P<num>[12])\.(?:f(ast)?q)(?:\.gz)?$", re.IGNORECASE),
]

def parse_args():
    ap = argparse.ArgumentParser(
        description="Scan an input directory and generate samples.tsv with columns: sample\\tR1\\tR2"
    )
    ap.add_argument("input_dir", type=Path, help="Directory containing FASTQ files")
    ap.add_argument("-o", "--output", type=Path, default=Path("samples.tsv"), help="Output TSV path")
    ap.add_argument("--recursive", action="store_true", help="Recursively search subdirectories")
    ap.add_argument("--sep", default="\t", help="Delimiter for TSV (default: TAB)")
    return ap.parse_args()

def infer_pair_key(name: str):
    """从文件名里识别 R1/R2/1/2，并返回 (prefix, readLabel)"""
    for pat in PATTERNS:
        m = pat.match(name)
        if m:
            gd = m.groupdict()
            prefix = gd.get("prefix")
            read = gd.get("read")
            num = gd.get("num")
            if read is None and num is not None:
                read = f"R{num}"
            return prefix, read.upper()
    return None, None

def main():
    args = parse_args()
    if not args.input_dir.exists() or not args.input_dir.is_dir():
        sys.exit(f"[Error] Input directory not found or not a directory: {args.input_dir}")

    # 收集 fastq/fq(.gz)
    files = []
    it = args.input_dir.rglob("*") if args.recursive else args.input_dir.glob("*")
    for p in it:
        if not p.is_file():
            continue
        low = p.name.lower()
        if low.endswith((".fq", ".fastq", ".fq.gz", ".fastq.gz")):
            files.append(p)

    if not files:
        sys.exit("[Error] No FASTQ files found in input_dir.")

    # 解析前缀并配对
    pairs = {}
    for p in files:
        prefix, read = infer_pair_key(p.name)
        if prefix is None or read not in {"R1", "R2"}:
            print(f"[Warn] Skip unrecognized filename: {p.name}", file=sys.stderr)
            continue
        pairs.setdefault(prefix, {"R1": None, "R2": None})
        if pairs[prefix][read] is not None:
            print(f"[Warn] Duplicate {read} for prefix {prefix}, keep first: {pairs[prefix][read]}", file=sys.stderr)
            continue
        pairs[prefix][read] = str(p.resolve())

    # 整理为 DataFrame
    rows = []
    for prefix, d in sorted(pairs.items()):
        rows.append((prefix, d.get("R1", ""), d.get("R2", "")))

    if not rows:
        sys.exit("[Error] No valid pairs constructed.")

    df = pd.DataFrame(rows, columns=["sample", "R1", "R2"])

    orphan = df[(df["R1"] == "") | (df["R2"] == "")]
    if len(orphan) > 0:
        print("[Warn] Found unpaired entries:", file=sys.stderr)
        for _, r in orphan.iterrows():
            print(f"  - {r['sample']}: R1={r['R1']}, R2={r['R2']}", file=sys.stderr)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(args.output, sep=args.sep, index=False)
    print(f"[OK] Wrote {len(df)} rows to {args.output}")

if __name__ == "__main__":
    sys.exit(main())