read fastq go

readfq.closure.go

package main

import (
	"bufio"
	"bytes"
	"compress/bzip2"
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"runtime/pprof"
	"strings"
)

func Xopen(fName string) (*os.File, *bufio.Reader) {
	var (
		fos *os.File
		err error
	)

	// Open the file (unless we are reading from stdin
	if fName != "-" {
		fos, err = os.Open(fName)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Can't open %s: error: %s\n", fName, err)
			os.Exit(1)
		}
	} else {
		fos = os.Stdin
	}

	// Deal with the compression of the file
	var newReader io.Reader
	err = nil
	if strings.HasSuffix(fName, ".gz") {
		newReader, err = gzip.NewReader(fos)
	} else if strings.HasSuffix(fName, ".bz2") {
		newReader = bzip2.NewReader(fos)
	} else {
		newReader = fos // No compression, stdin
	}

	if err != nil {
		fmt.Fprintf(os.Stderr,
			"%s file is not in the correct format: error: %s\n", fName, err)
		os.Exit(1)
	}

	// We want a buffered stream to increase performance
	return fos, bufio.NewReader(newReader)
}

type record struct {
	name, seq, qual string
}

type fqReader struct {
	r *bufio.Reader
	last, seq, qual []byte
	finished bool
	rec record
}

func (fq *fqReader) iterLines() ([]byte, bool) {
	line, err := fq.r.ReadSlice('\n')
	if err != nil {
		if err == io.EOF {
			return line, true
		} else {
			panic(err)
		}
	}
	return line, false
}

var space = []byte(" ")

func (fq *fqReader) iter() (record, bool) {
	if fq.finished {
		return fq.rec, fq.finished
	}
	// Read the seq id (fasta or fastq)
	if fq.last == nil {
		for l, done := fq.iterLines(); !done; l, done = fq.iterLines() {
			if l[0] == '>' || l[0] == '@' { // read id
				fq.last = l[0 : len(l)-1]
				break
			}
		}
		if fq.last == nil { // We couldn't find a valid record, no more data in file
			fq.finished = true
			return fq.rec, fq.finished
		}
	}
	fq.rec.name = string(bytes.SplitN(fq.last, space, 1)[0])
	fq.last = nil

	// Now read the sequence
	fq.seq = fq.seq[:0]
	for l, done := fq.iterLines(); !done; l, done = fq.iterLines() {
		c := l[0]
		if c == '+' || c == '>' || c == '@' {
			fq.last = l[0 : len(l)-1]
			break
		}
		fq.seq = append(fq.seq, l[0:len(l)-1]...)
	}
	fq.rec.seq = string(fq.seq)

	if fq.last == nil { // Reach EOF
		fq.finished = true
	}
	if fq.last[0] != '+' { // fasta record; set sequence
		return fq.rec, fq.finished
	}
	leng := 0
	fq.qual = fq.qual[:0]
	for l, done := fq.iterLines(); !done; l, done = fq.iterLines() {
		fq.qual = append(fq.qual, l[0:len(l)-1]...)
		leng += len(l)
		if leng >= len(fq.seq) { // we have read enough quality
			fq.last = nil
			fq.rec.qual = string(fq.qual)
			return fq.rec, fq.finished
		}
	}
	fq.finished = true
	fq.rec.qual = string(fq.qual)
	return fq.rec, fq.finished // incomplete fastq quality, return what we have
}

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

func main() {
	flag.Parse()
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	fp, r := Xopen("-")
	defer fp.Close()

	n, sLen, qLen := 0, int64(0), int64(0)
	var fqr fqReader
	fqr.r = r
	for r, done := fqr.iter(); !done; r, done = fqr.iter() {
		n += 1
		sLen += int64(len(r.seq))
		qLen += int64(len(r.qual))
	}
	fmt.Println(n, "\t", sLen, "\t", qLen)
}