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 
  32     @staticmethod
  33     def make_dna(fasta_str):
  34         lines = fasta_str.splitlines()
  35         title = lines[0][1:]
  36         sequence = []
  37         for line in lines[1:]:
  38             line = line.strip()
  39             sequence.append(line)
  40         sequence = ''.join(sequence)
  41         return DNA(title, sequence)
  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 
  55 def main(input_file, output_file):
  56     for fasta in FastaIterator(input_file):
  57         dna = DNA.make_dna(fasta)
  58         rc_dna = dna.get_reverse_complement()
  59         output_file.write(rc_dna.write())
  60     
  61 
  62 class DNATest(unittest.TestCase):
  63     def test_dna_creation(self):
  64         dna = DNA('test1', 'AGTCAA')
  65         self.assertEquals('test1', dna.title)
  66         self.assertEquals('AGTCAA', dna.sequence)
  67 
  68     def test_get_reverse_complement(self):
  69         dna = DNA('test1', 'AGTCAA')
  70         rc_dna = dna.get_reverse_complement()
  71         self.assertEquals('test1 reverse complement', rc_dna.title)
  72         self.assertEquals('TTGACT', rc_dna.sequence)
  73 
  74     def test_make_dna_from_fasta(self):
  75         input = """\
  76 >test1
  77 AGTC
  78 AGTC
  79 """
  80         dna = DNA.make_dna(input)
  81         self.assertEquals('test1', dna.title)
  82         self.assertEquals('AGTCAGTC', dna.sequence)
  83 
  84     def test_fasta_iterator(self):
  85         input = StringIO("""\
  86 >test1
  87 AGTC
  88 AGTC
  89 >test2
  90 AGTC
  91 AGTA
  92 >test3
  93 AGTG
  94 AGTC
  95 """)
  96         expect1 = """\
  97 >test1
  98 AGTC
  99 AGTC
 100 """
 101         expect2 = """\
 102 >test2
 103 AGTC
 104 AGTA
 105 """
 106         fi = FastaIterator(input)
 107         self.assertEquals(expect1, fi.next())
 108         self.assertEquals(expect2, fi.next())
 109 
 110     def test_write_dna(self):
 111         dna = DNA('test1', 'AGTCAAGG')
 112         expect = """\
 113 >test1
 114 AGTCA
 115 AGG
 116 """
 117         self.assertEquals(expect, dna.write(cols=5))
 118         
 119 
 120 if __name__ == '__main__':
 121     #unittest.main()
 122     import sys
 123     main(sys.stdin, sys.stdout)