printing the sum of the counts of a string over multiple lines in Python

Question

This is a simple question but I just can't solve it. I want to count the number of A for each sequence of lines. Please see the example below:

This is my input:

>sca200
ACACGTGYNNNN
ACGTCCCGWCNN
NNNNNNNNNA
>scaf500
AAAAAAAAAAAA
TTTTTTTTTTTT
WCWCWNNNN
>scaf201
AACACACACACC
GTGTGTGTGTGT
WWRRRYNNNNNN
NNNNNN

code:

#!/usr/bin/python
from __future__ import division
import sys

fasta = open(sys.argv[1], "r")

for line in fasta:
    line = line.rstrip("\n")
    if line.startswith(">"):
        total_A = 0
        print line[1:]
    else:
        A = line.count('A')
        total_A = total_A + A
        print total_A

The output is:

sca200
2
3
4
scaf500
12
12
12
scaf201
6
6
6
6

How can I get it to report only the final number?, that is:

sca200
4
scaf500
12
scaf201
6

Answer 1

try to split input by newlines and then output totals only if the line starts with < (what does mean that sequence is over) - remember about printing first and last record out of the loop

#!/usr/bin/python
from __future__ import division
import sys

#reading file into str variable
fasta_file = open(sys.argv[1], "r")
fasta = fasta_file.read()
fasta_file.close()

total = 0

print fasta.split('\n')[0][1:]

for f in fasta.split('\n')[1:]:
    if f[0] != '>':
        total += f.count('A')
    else:
        print total, '\n', f[1:]
        total = 0

print total

Answer 2

This should solve your problem:

#!/usr/bin/python
from __future__ import division
import sys

fasta = open(sys.argv[1], "r")
total_A = None
for line in fasta:
    line = line.rstrip("\n")
    if line.startswith(">"): 
        print total_A if total_A != None else 0
        total_A = 0
        print line[1:]
    else:
        A = line.count('A')
        total_A += A
 print total_A

You just want to print the total count of A just when a new fasta header starts.

Note: Edited to address a comment raised by @Lafexlos.

Answer 3

Here is one-liner solution:

from __future__ import print_function

import sys
import re

with open(sys.argv[1], 'r') as f:
    data = f.read()

"""
1.  find all blocks of text and split it into two groups: (block_name, corresponding_TEXT)
2.  loop through blocks
3.  print 'block_name' and the length of list containing all 'A's from the corresponding_TEXT
"""

[   print('{0}\n{1}'.format(name, len(re.findall(r'A', txt, re.M)))) 
    for name, txt in re.findall(r'>(sca[^\n]*)([^>]*)', data, re.M)
]

Output:

sca200
4
scaf500
12
scaf201
6