1 """ReverseComplement using Generator by yong27, 2005-05-18
   2 """
   3 import os, unittest
   4 from cStringIO import StringIO
   5 
   6 class Fasta:
   7     def __init__(self):
   8         self.title = ''
   9         self.sequence = ''
  10     def setSequence(self, aSequence):
  11         self.sequence = aSequence
  12     def getRcSequence(self):
  13         DnaComp = {'A':'T','C':'G','G':'C','T':'A'}
  14         result = []
  15         listSeq = list(self.sequence)
  16         listSeq.reverse()
  17         for c in listSeq:
  18             result.append(DnaComp[c])
  19         return ''.join(result)
  20     def setTitle(self, aTitle):
  21         self.title = aTitle
  22     def getTitle(self):
  23         return self.title
  24     def getSequence(self):
  25         return self.sequence
  26 
  27 class FastaGenerator:
  28     def getFasta(self, aStream):
  29         f= Fasta()
  30         listSeq = []
  31         for line in aStream.splitlines():
  32             if line[0] == '>':
  33                 f.setTitle(line[1:])
  34             else:
  35                 listSeq.append(line)
  36         f.setSequence(''.join(listSeq))
  37         return f
  38 
  39     def getFastaByLines(self, lines):
  40         return self.getFasta(os.linesep.join(lines))
  41 
  42     def parse(self, aFile):
  43         lines = [aFile.next().strip()]
  44         for line in aFile:
  45             if line[0] == '>':
  46                 yield self.getFastaByLines(lines)
  47                 lines = []
  48             lines.append(line.strip())
  49         else:
  50             yield self.getFastaByLines(lines)
  51 
  52 
  53 class FastaTest(unittest.TestCase):
  54     def setUp(self):
  55         self.f = Fasta()
  56     def testReverseComplement(self):
  57         self.f.setSequence('AGTC')
  58         self.assertEquals('GACT', self.f.getRcSequence())
  59 
  60     def testGetTitle(self):
  61         self.f.setTitle('aTitle')
  62         self.f.setSequence('AGTC')
  63         self.assertEquals('aTitle', self.f.getTitle())
  64         self.assertEquals('AGTC', self.f.getSequence())
  65 
  66 
  67 class GeneratorTest(unittest.TestCase):
  68     def setUp(self):
  69         self.g = FastaGenerator()
  70         self.single = """\
  71 >a
  72 AGTC
  73 GTC
  74 """
  75         self.multiple = """\
  76 >a
  77 AGTC
  78 GTC
  79 >b
  80 AAA
  81 GG 
  82 """
  83     def testSingle(self):
  84         f = self.g.getFasta(self.single)
  85         self.assertEquals('a', f.getTitle())
  86         self.assertEquals('AGTCGTC', f.getSequence())
  87 
  88     def testMultiple(self):
  89         g = self.g.parse(StringIO(self.multiple))
  90         f = g.next()
  91         self.assertEquals('a', f.getTitle())
  92         self.assertEquals('AGTCGTC', f.getSequence())
  93         f = g.next()
  94         self.assertEquals('b', f.getTitle())
  95         self.assertEquals('AAAGG', f.getSequence())
  96 
  97 if __name__=='__main__':
  98     unittest.main()