countReads

#! /usr/bin/env python3

import os
import sys
import signal
import glob
import logging
import argparse
from multiprocessing import Pool

import pysam

parser = argparse.ArgumentParser()
parser.add_argument(dest="input", help="BAM files or Directory with BAM files.", nargs="+")
parser.add_argument('-p', '--processes', help="Number of processes to use.", default=1, type=int)
parser.add_argument('-f', '--include-flags', help="SAM flags to include (default: 3)", default=3, type=int)
parser.add_argument('-F', '--exclude-flags', help="SAM flags to exclude (default: 3340)", default=3340, type=int)
parser.add_argument('-d', '--dry-run', help="Dry run", action="store_true")

args = parser.parse_args()

def worker_init():
    signal.signal(signal.SIGINT, signal.SIG_IGN)

def get_total_reads(bam):
    num_reads = int(pysam.view("-c", f"-f {args.include_flags}", f"-F {args.exclude_flags}", bam).strip())
    print(f"Found {num_reads} reads in {os.path.basename(bam)}.", file=sys.stderr)

    return {bam: num_reads}

if __name__=="__main__":
    all_bams = []
    for path in args.input:
        if os.path.isdir(path):
            all_bams.extend(glob.glob(f"{path}/**/*.bam", recursive=True))
        elif os.path.isfile(path):
            all_bams.extend([path])

    print(f"Found {len(all_bams)} BAM file(s).", file=sys.stderr)
    if len(all_bams) < 1:
        print("done.", file=sys.stderr)
        sys.exit(0)

    if args.dry_run:
        print("\nWill be counting following files:\n---------------------------------")
        for file in all_bams:
            print(file)
        sys.exit(0)

    # Collect read metrics for all input files
    logging.info("Computing total reads in each file..")
    pool = Pool(processes=args.processes, initializer=worker_init)
    async_result = pool.map_async(get_total_reads, all_bams)

    number_reads = {}
    try:
        result = async_result.get()
        for k in result:
            number_reads.update(k)
    except Exception as e:
        logging.error(e)
    finally:
        pool.close()

    for k, v in number_reads.items():
        sys.stdout.writelines(f"{k}\t{v}\n")