查找行，然后获得下一行

Question

I have the following problem: I open a file and read it line by line searching for a specific pattern. When I found it, I would like to write the entire line AND THE NEXT TWO LINES into a new file. The problem is that I don't know how to get from the line that I've found to the next 2.

AAA
XXX
XXX
BBB
XXX
XXX
CCC
XXX
XXX

In this example it would be that I find "BBB" and I want to get the next two lines.

What could I do? Thank you very much for your kind help!

Edit: I realized that I have to ask more precisely.

This is the code from my colleague

for k in range(0,len(watcrd)):
    if cvt[k]>cvmin:
        intwat+=1
        sumcv+=cvt[k]
        sumtrj+=trj[k]/((i+1)*sep/100)
        endline='%5.2f %5.2f' % (cvt[k],trj[k]/((i+1)*sep/100)) # ivan
        ftrj.write(watline[k][:55]+endline+'\n')
        fall.write(watline[k][:55]+endline+'\n')

For every k in range I would like to write k, k+1, k+2 to the file ftrj. Which is the best way to do this?

Edit 2: I am sorry, but I realized that I've made a mistake. What you suggested worked, but I realized that I have to include it in a different part of the code.

for line in lines[model[i]:model[i+1]]:
    if line.startswith('ATOM'):
    resi=line[22:26]
    resn=line[17:20]
    atn=line[12:16]
    crd=[float(line[31:38]),float(line[38:46]),float(line[46:54])]

    if (resn in noprot)==False and atn.strip().startswith('CA')==True:
        protcrd.append(crd)
    if (resn in reswat)==True and (atn.strip() in atwat)==True:
        watcrd.append(crd)
        watline.append(line)

I would think of something like this:

    (...)
    if (resn in reswat)==True and (atn.strip() in atwat)==True:
        watcrd.append(crd)
        watline.append(line)
        for i in range(1, 3):
            try:
                watcrd.append(crd[line + i])
                watline.append(line[line + i])
            except IndexError:
                break

But it doesn't work. How can I indicate the part and the line that I want to append to this list?

Answer 1

Python file objects are iterators, you can always ask for the next lines:

with open(inputfilename) as infh:
    for line in infh:
        if line.strip() == 'BBB':
            # Get next to lines:
            print next(infh)
            print next(infh)

Here next() function advances the infh iterator to the next line, returning that line.

However, you are not processing a file; you are processing a list instead; you can always access later indices in the list:

    ftrj.write(watline[k][:55]+endline+'\n')
    fall.write(watline[k][:55]+endline+'\n')
    for i in range(1, 3):
        try:
            ftrj.write(watline[k + i][:55]+endline+'\n')
            fall.write(watline[k + i][:55]+endline+'\n')
        except IndexError:
            # we ran out of lines in watline
            break