1 import unittest
   2 from cStringIO import StringIO
   3 
   4 class DNA:
   5     def __init__(self, title, sequence):
   6         self.title = title
   7         self.sequence = sequence
   8 
   9     def get_reverse_complement(self):
  10         base_pair = {
  11             'A': 'T',
  12             'T': 'A',
  13             'G': 'C',
  14             'C': 'G',
  15         }
  16         title = "%s reverse complement" % self.title
  17         new_sequence = []
  18         for nucleotide in self.sequence:
  19             new_sequence.append(base_pair[nucleotide])
  20         new_sequence = ''.join(reversed(new_sequence))
  21         return DNA(title, new_sequence)
  22 
  23     def write(self, cols=60):
  24         result = [">%s" % self.title]
  25         sequence = self.sequence[:]
  26         while sequence:
  27             result.append(sequence[:cols])
  28             sequence = sequence[cols:]
  29         return "\n".join(result) + '\n'
  30         
  31     @staticmethod
  32     def make_dna(fasta_str):
  33         lines = fasta_str.splitlines()
  34         title = lines[0][1:]
  35         sequence = []
  36         for line in lines[1:]:
  37             line = line.strip()
  38             sequence.append(line)
  39         sequence = ''.join(sequence)
  40         return DNA(title, sequence)
  41 
  42 
  43 def FastaIterator(afile):
  44     lines = [afile.readline()]
  45     for line in afile:
  46         if line.startswith('>'):
  47             yield ''.join(lines)
  48             lines = [line]
  49         else:
  50             lines.append(line)
  51     yield ''.join(lines)
  52 
  53 
  54 def main(input_file, output_file):
  55     for fasta in FastaIterator(input_file):
  56         dna = DNA.make_dna(fasta)
  57         rc_dna = dna.get_reverse_complement()
  58         output_file.write(rc_dna.write())
  59     
  60 
  61 class DNATest(unittest.TestCase):
  62     def test_dna_creation(self):
  63         dna = DNA('test1', 'AGTCAA')
  64         self.assertEquals('test1', dna.title)
  65         self.assertEquals('AGTCAA', dna.sequence)
  66 
  67     def test_get_reverse_complement(self):
  68         dna = DNA('test1', 'AGTCAA')
  69         rc_dna = dna.get_reverse_complement()
  70         self.assertEquals('test1 reverse complement', rc_dna.title)
  71         self.assertEquals('TTGACT', rc_dna.sequence)
  72 
  73     def test_make_dna_from_fasta(self):
  74         input = """\
  75 >test1
  76 AGTC
  77 AGTC
  78 """
  79         dna = DNA.make_dna(input)
  80         self.assertEquals('test1', dna.title)
  81         self.assertEquals('AGTCAGTC', dna.sequence)
  82 
  83     def test_fasta_iterator(self):
  84         input = StringIO("""\
  85 >test1
  86 AGTC
  87 AGTC
  88 >test2
  89 AGTC
  90 AGTA
  91 >test3
  92 AGTG
  93 AGTC
  94 """)
  95         expect1 = """\
  96 >test1
  97 AGTC
  98 AGTC
  99 """
 100         expect2 = """\
 101 >test2
 102 AGTC
 103 AGTA
 104 """
 105         fi = FastaIterator(input)
 106         self.assertEquals(expect1, fi.next())
 107         self.assertEquals(expect2, fi.next())
 108 
 109     def test_write_dna(self):
 110         dna = DNA('test1', 'AGTCAAGG')
 111         expect = """\
 112 >test1
 113 AGTCA
 114 AGG
 115 """
 116         self.assertEquals(expect, dna.write(cols=5))
 117         
 118 
 119 if __name__ == '__main__':
 120     #unittest.main()
 121     import sys
 122     main(sys.stdin, sys.stdout)