bam_region_extract.py

#!/usr/bin/env python

from path import Path
import sh
import argparse
from functools import partial
import string

def parse_args():
    parser = argparse.ArgumentParser(
        description='extract regions from bam into lots of bams',
    )
    parser.add_argument(
        'bamfile',
        type=Path,
        help='Bam file to extract from'
    )
    parser.add_argument(
        'regions',
        type=Path,
        help='File containing region strings to extract from bam'
    )
    parser.add_argument(
        '--width',
        type=int,
        default=100,
        help='How wide the region is to extract. So half this on the left and half on the right of the given position'
    )
    parser.add_argument(
        '--outdir',
        default=Path('.'),
        type=Path,
        help='path to put each extracted region\'s bam output into'
    )
    return parser.parse_args()

def make_region_bam(inputbam, regionstr, outputbam):
    _c = sh.samtools.bake('sort')
    _c(sh.samtools('view', inputbam, regionstr, h=True, _piped=True), _out=outputbam)
    sh.samtools('index', outputbam)

def make_regionstr(ref_pos, width):
    '''
    >>> make_regionstr('FOO:500', 100)
    'FOO:450-550'
    >>> make_regionstr('FOO:1', 100)
    'FOO:1-51'
    '''
    ref, pos = ref_pos.split(':')
    pos = int(pos)
    delta = int(width / 2.0)
    lpos = max(1, pos - delta)
    rstring = '{0}:{1}-{2}'.format(ref, lpos, pos+delta)
    return rstring

def main():
    args = parse_args()
    args.outdir.makedirs_p()
    rstr = partial(make_regionstr, width=args.width)
    regions = map(rstr, args.regions.lines())
    for regionstr in regions:
        outbam, _ = args.bamfile.basename().splitext()
        outbam = args.outdir / outbam + '_{0}.bam'.format(regionstr.replace('/',''))
        print "Extracting region: {0} into {1}".format(regionstr, outbam)
        make_region_bam(args.bamfile, regionstr, outbam)

if __name__ == '__main__':
    main()

great - can we add this to bio-bits?