add counts of variation samples to annotated VCF file

add_count_samples.py

#!/usr/bin/env python
#!-*- coding: utf-8 -*-

import sys


def add_count_samples(infile, outfile=None):
    outfile = outfile or infile+'.addCountSamples'
    with open(infile) as f, open(outfile, 'w') as out:
        for line in f:
            if line.startswith('Priority'):
                title = line.strip().split('\t')
                priority_index = title.index('Priority')
                format_index = title.index('FORMAT')
                ori_ref_index = title.index('Ori_REF')
                samples = title[format_index+1:ori_ref_index]

                title[priority_index+1:priority_index+1] = ['Count', 'Samples']
                out.write( '\t'.join(title)+'\n' )

                continue

            linelist = line.strip().split('\t')
            infos = linelist[format_index+1:ori_ref_index]
            count = 0
            sample_alts = []
            for sample, info in zip(samples, infos):
                if info != '.':
                    count += 1
                    sample_alts.append(sample)
            linelist[priority_index+1:priority_index+1] = [str(count), ','.join(sample_alts)]
            out.write( '\t'.join(linelist)+'\n' )



if __name__ == '__main__':
    if len(sys.argv) < 2:
        print "Usage: python %s <infile> [outfile]" % sys.argv[0]
        exit(1)
    add_count_samples(*sys.argv[1:])