read mixed(or not) fasta/fastq files

readfq.object.go

package main

import (
  "bufio"
  "bytes"
  "flag"
  "fmt"
  "github.com/drio/drio.go/common/files"
  "io"
  "log"
  "os"
  "runtime/pprof"
)

type record struct {
  name, seq, qual string
}

type fqReader struct {
  r               *bufio.Reader
  last, seq, qual []byte // last line processed, temporary seq and qual values
  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 { // There are more lines
    if fq.last[0] != '+' { // fasta record
      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 := files.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)
}