#format python
"""ReverseComplement using Generator by yong27, 2005-05-18
"""
import os, unittest
from cStringIO import StringIO

class Fasta:
    def __init__(self):
        self.title = ''
        self.sequence = ''
    def setSequence(self, aSequence):
        self.sequence = aSequence
    def getRcSequence(self):
        DnaComp = {'A':'T','C':'G','G':'C','T':'A'}
        result = []
        listSeq = list(self.sequence)
        listSeq.reverse()
        for c in listSeq:
            result.append(DnaComp[c])
        return ''.join(result)
    def setTitle(self, aTitle):
        self.title = aTitle
    def getTitle(self):
        return self.title
    def getSequence(self):
        return self.sequence

class FastaGenerator:
    def getFasta(self, aStream):
        f= Fasta()
        listSeq = []
        for line in aStream.splitlines():
            if line[0] == '>':
                f.setTitle(line[1:])
            else:
                listSeq.append(line)
        f.setSequence(''.join(listSeq))
        return f

    def getFastaByLines(self, lines):
        return self.getFasta(os.linesep.join(lines))

    def parse(self, aFile):
        lines = [aFile.next().strip()]
        for line in aFile:
            if line[0] == '>':
                yield self.getFastaByLines(lines)
                lines = []
            lines.append(line.strip())
        else:
            yield self.getFastaByLines(lines)


class FastaTest(unittest.TestCase):
    def setUp(self):
        self.f = Fasta()
    def testReverseComplement(self):
        self.f.setSequence('AGTC')
        self.assertEquals('GACT', self.f.getRcSequence())

    def testGetTitle(self):
        self.f.setTitle('aTitle')
        self.f.setSequence('AGTC')
        self.assertEquals('aTitle', self.f.getTitle())
        self.assertEquals('AGTC', self.f.getSequence())


class GeneratorTest(unittest.TestCase):
    def setUp(self):
        self.g = FastaGenerator()
        self.single = """\
>a
AGTC
GTC
"""
        self.multiple = """\
>a
AGTC
GTC
>b
AAA
GG 
"""
    def testSingle(self):
        f = self.g.getFasta(self.single)
        self.assertEquals('a', f.getTitle())
        self.assertEquals('AGTCGTC', f.getSequence())

    def testMultiple(self):
        g = self.g.parse(StringIO(self.multiple))
        f = g.next()
        self.assertEquals('a', f.getTitle())
        self.assertEquals('AGTCGTC', f.getSequence())
        f = g.next()
        self.assertEquals('b', f.getTitle())
        self.assertEquals('AAAGG', f.getSequence())

if __name__=='__main__':
    unittest.main()