Split a FASTQ file by header

split_by_header.py

#!/usr/bin/python
"""Split up a FASTQ file based on the first field of the header."""

from collections import defaultdict
import gzip
import sys
import os

fp = sys.argv[1]
if not os.path.exists(fp):
    raise Exception("{} does not exist".format(fp))

folder = "{}.split".format(fp)

if not os.path.exists(folder):
    os.mkdir(folder)

if fp.endswith('.gz'):
    f = gzip.open(fp)
else:
    f = open(fp)

record = []
records = defaultdict(list)
for line in f:
    name = line[1:].split('_')[0]
    records[name].append(line)
    records[name].append(f.readline())
    records[name].append(f.readline())
    records[name].append(f.readline())
    if len(records[name]) > 100000:
        # print "Writing {} lines".format(len(records[name]))
        with open("{}/{}.fastq".format(folder, name), 'a') as fo:
            fo.write(''.join(records[name]))
        records[name] = []

f.close()

for name in records.keys():
    with open("{}/{}.fastq".format(folder, name), 'a') as fo:
            fo.write(''.join(records[name]))