Problem in Python with function only writing one line to file

I am 开发者_StackOverflowtrying to write a function in Python 3 that will write all lines that end with the string 'halloween' to a file. When I call this function, I can only get one line to write to the output file (file_2.txt). Can anyone point out where my problem is? Thanks in advance.

def parser(reader_o, infile_object, outfile_object):
    for line in reader_o:
        if line.endswith('halloween'):
            return(line)

with open("file_1.txt", "r") as file_input:
    reader = file_input.readlines()
    with open("file_2.txt", "w") as file_output:
        file_output.write(parser(reader))

---

def parser(reader_o):
    for line in reader_o:
        if line.rstrip().endswith('halloween'):
            yield line

with open("file_1.txt", "r") as file_input:
    with open("file_2.txt", "w") as file_output:
        file_output.writelines(parser(file_input))

This is called a generator. It can also be written as an expression instead of a function:

with open("file_1.txt", "r") as file_input:
    with open("file_2.txt", "w") as file_output:
        file_output.writelines(line for line in file_input if line.rstrip().endswith('halloween'))

If you're on Python 2.7 / 3.2, you can do the two withs like this:

with open("file_1.txt", "r") as file_input, open("file_2.txt", "w") as file_output:

You don't need to do readlines() on the file, just telling the loop to iterate over the open file itself will do the exact same thing.

Your problem was that return always would exit the loop on the first match. yield stops the loop, passes out the value, then the generator can be started again from the same point.

---

line.endswith('halloween') might only work on the last of a file, since all other lines will have a newline appended. rstrip the line first. Also, use yield instead of return.

if line.rstrip().endswith('halloween'):
    yield line

Note that this will also strip off spaces at the end of the line, which may or may not be what you want.

You'll also have to modify your consumer to

with open("file_2.txt", "w") as file_output:
    for ln in parser(reader):
        file_output.write(ln)

---

Perhaps your parser function should be a generator. At the moment it's only called once and returns the first line that has "halloween" in it.

Like the following:

def parser(reader_o):
    for line in reader_o:
        if line.endswith('halloween'):
            yield line

with open("file_1.txt", "r") as file_input:
    with open("file_2.txt", "w") as file_output:
        file_output.writelines(parser(file_input))