Multiple Testcases Stdin (Python)

Question

The first question of this year's facebook hacker cup had input of the following form:

 3 #number of test cases
 4 #number of rows of test case 1 
 . . . x
 . . x x
 . . x x
 . . . x
 2 #number of rows of test case 2
 . . x x
 . . x x
 3 #number of rows of test case 3 
 x x . .
 x x . .
 . . . x

Normally when doing codeforces problems or topcoder, you don't have to input 5 test cases after each other, you just make it for one, and they run it through 20-25 testcases.

I had a lot of struggle trying to manipulate this data to make it usable and was wondering how one might do it.

For example, if it was just

 5
 2 3 4 5 6

I could use input() to obtain the first number, and

import sys
data = []
for line in sys.stdin:
    y = [int(x) for x in line.split()]
    data.append(y)

to manipulate the rest. If I did something like this for this problem (replacing int with str), I would end up with one array like [3,4,data,2,data,3,data] which seems difficult to manipulate.

How can I read multiple test cases from stdin? (even general answers helpful since the question itself is not that specific)

Answer 1

I tend to wrap this up in a generator. For example:

import sys

def read_data(source):
    N = int(next(source))
    for case in range(N):
        num_rows = int(next(source))
        rows = [next(source).split() for i in range(num_rows)]
        yield rows

for case in read_data(sys.stdin):
    print case

produces

dsm@notebook:~/coding$ cat source.txt | python getdata.py 
[['.', '.', '.', 'x'], ['.', '.', 'x', 'x'], ['.', '.', 'x', 'x'], ['.', '.', '.', 'x']]
[['.', '.', 'x', 'x'], ['.', '.', 'x', 'x']]
[['x', 'x', '.', '.'], ['x', 'x', '.', '.'], ['.', '.', '.', 'x']]

This way, the data reader doesn't care if the source is stdin, or a file, or whatever, and you could pass it something which stripped comments if necessary.