add_fragment_to_bam.py

#!/bin/env python

import os
import sys
import getopt
import pysam


def bed_reader(hin=sys.stdin):
    """
    Reads a bed file
    """
    # foreach line in the input stream
    for line in hin:

        # get the next line
        line = line.rstrip()
        if len(line) == 0:
            continue

        # get the fields in the table
        fields = line.split("\t")
        if len(fields) < 3:
            continue

        # parse the start and end
        try:
            start = int(fields[1])
            end = int(fields[2])
        except ValueError:
            continue

        # check for additional columns
        add = []
        if len(fields) > 3:
            add = fields[3:]

        # return the data columns
        yield fields[0], start, end, add


def main(fn_sam_in=None, fn_sam_out=None, bedfile=sys.stdin, tagname="fr"):
    """
    The main program loop
    """
    # open the source data
    infile = pysam.AlignmentFile(fn_sam_in, 'rb')
    outfile = pysam.AlignmentFile(fn_sam_out, "wb", template=infile)

    # for each bed region
    for sname, start, end, other in bed_reader(bedfile):
        tag = "%s:%d-%d:%s" % (sname, start, end, "_".join(other))
        for read in infile.fetch(sname, start, end):
            read.set_tag(tagname, tag, "Z")
            outfile.write(read)

    # close the input and output files
    infile.close()
    outfile.close()

if __name__ == "__main__":


    def usage(message="", error=None):
        """
        Print the usage information
        """
        sys.stdout.write("""
Message %(message)s

Called as %(commandline)s

Usage:
    python %(tool)s -i [input file] -b [bed file] -o [output file] 

Options:
    -i/--input    [file]  the input BAM file
    -b/--bed      [file]  a bed file with restriction sites, default stdin
    -o/--output   [file]  the output BAM file 
    -h/--help       print this message

""" % {
    "message": message,
    "tool": sys.argv[0],
    "commandline": " ".join(sys.argv)
})
        if error is not None:
            sys.exit(error)
        # end of usage

    # main variables
    fname_input = None
    fname_output = None
    h_bed = sys.stdin

    # option parsing
    shortopt = "i:o:b:h"
    longopt = ["input=", "output=", "bed=", "help"]

    try:
        opts, args = getopt.getopt(sys.argv[1:], shortopt, longopt)
        for o, a in opts:
            if o in ("-i", "--input"):
                fname_input = a
            elif o in ("-o", "--output"):
                fname_output = a
            elif o in ("-b", "--bed"):
                h_bed = open(a, "r")
            elif o in ('-h', '--help'):
                usage("Help was asked", error=0)
    except getopt.GetoptError as err:
        usage(str(err), error=2)

    if fname_input is None or not os.path.exists(fname_input):
        usage("Input file %s not found" % str(fname_input), error=101)

    if fname_output is None:
        usage("Output file %s not specified" % str(fname_output), error=102)

    main(fname_input, fname_output, h_bed)