使用生成器过滤具有多行记录的文件的python3

Question

I need to read huge files structured as multiline records and write to a file records with certain indices, say record numbers R = 1, 2 and 1093. If records are N = 3 lines each this amounts to reading the file line by line and then write lines numbers 1, 2, 3 and 4, 5, 6 and 3277, 3278, 3279 (in that the first line in each record Ri starts at line number Ri-1 * N + 1.

I guess one could calculate the lines to write and go through the file line by line and write those lines. However, is it possible to "zip" consecutive lines 1, 2 and 3 into a generator object containing records and filter these somehow or would or print them directly to file if they enumerate to R ? Something along this pseudocode :

def subset(file_in, file_out, N, R):
    with open(file_in, "rt") as fin, open(file_out, "wt") as fout:
        line = (line.rstrip() for line in fin)
        record = enumerate(zip(line, line, line)) # What if records are of size N
        for i, r in record if i in R:
            fout.write(r)

What to do if you want the record size N as parameter ?

UPDATE EXAMPLE

An example for file_in (4 records, 3 lines/record):

dslfkj
2
a
dflkj
3
g
fds
2
b
fsdlkj
1
n

Then subset(file_in, file_out, 3, [1,3]) would give (file_out)

dslfkj
2
a
fds
2
b

Answer 1

For this problem, it makes sense just to tackle this directly line by line, using floor division.

For Example:

fin = '''
dslfkj
2
a
dflkj
3
g
fds
2
b
fsdlkj
1
'''

line_gen = (line.rstrip() for line in fin.strip().split())

R = [1, 3]
R = [val - 1 for val in R] #zero indexing
N = 3
for i, line in enumerate(line_gen):
    if i // N in R:
        print(line)

Output:

dslfkj
2
a
fds
2
b

Your function can look something like follows: (you may want to check if it works out of the box or requires tweaks. i did not check the file opening portion.

def subset(file_in, file_out, N, R):
    R = [val - 1 for val in R] #zero indexing
    with open(file_in, "rt") as fin, open(file_out, "wt") as fout:
        line_gen = (line.rstrip() for line in fin)
        for i, line in enumerate(line_gen):
            if i // N in R:
                fout.write(line)
                fout.write('\n')

---

Edit: The answer below pertains to how you could use the generators and group the values together. Having said that, i do not think you should need to use it. However, if you still wish to, you can construct your function based off of it.

Old answer:

You can create n references to the object using list, and then unpack using the * (aka splat) operator .

For example:

from itertools import zip_longest
line = (x for x in range(100, 132))
n = 3
record = zip(*([line] * n)) #equivalent to *[line, line, line] which is unpacked into zip arguments
for i, r in enumerate(record):
    print(i, r)

0 (100, 101, 102)
1 (103, 104, 105)
2 (106, 107, 108)
3 (109, 110, 111)
4 (112, 113, 114)
5 (115, 116, 117)
6 (118, 119, 120)
7 (121, 122, 123)
8 (124, 125, 126)
9 (127, 128, 129)

Also, depending on what you want to happen for "leftover" lines, you may want to use zip_longest instead.