在文件中将字符串匹配后保存n行到变量

Question

I'm trying to write a script that can read through a text file and save specific lines as lists. The text file looks something like this:

TRIAL1

Energy: 54432

Coordinates:

0.7 0.4 0.5

0.3 0.4 0.1

0.3 0.4 0.3

there are many more trials in the file (TRIAL2, TRIAL3 etc. with corresponding energy and coordinates), but I want to be able to select one particular trial and save its corresponding coordinates in a list (for later manipulation).

So far, I've written:

with open('energy.txt') as f:
   for line in f:
        if 'TRIAL1' in line:

I'm unsure how to then skip the lines "Energy" and "Coordinates" line, and save only the coordinates to a variable.

Answer 1

Assuming the "TRIAL" blocks all look the same as your example (adapt to fit your script):

lines = """TRIAL1
Energy: 54432
Coordinates:
0.7 0.4 0.5
0.3 0.4 0.1
0.3 0.4 0.3
TRIAL2
Energy: 54432
Coordinates:
test1
test2
test3"""

# initialize our control variables
counter = 0
collect = False

for line in lines.splitlines():
    if 'TRIAL2' in line:
        # mark the beginning of the block
        counter = 0
        collect = True
        continue
    if not collect:
        # nothing to do, go to next line
        continue
    else:
        # count the lines
        counter += 1
        if 2 < counter < 6:
            # use the 3rd, 4th and 5th line
            print(line)
        elif counter == 6:
            # stop collection at the 6th line of the block
            collect = False

Answer 2

I suggest making some assumptions about your input and just read each line that you expect to be there. You can later go back to add error checking code that enforces your assumptions. Something like this:

with open('energy.txt') as f:
    while True:
        trial = f.readline() # assume this line is "TRIAL#"
        energy = f.readline() # assume this line is "ENERGY:..."
        coordinates_label = f.readline() # assume this line is "Coordinates:"
        coordinates = f.readline()
        while "TRIAL" not in coordinates
            # parse the coordinates here

To ignore a line, just don't assign f.readline() to a variable. If there are really blank lines between each line of input, just add appropriate f.readline() calls.

Note: improving the while True: loop is left as an exercise for the reader.