Adding extra statements in a python list comprehension

Question

I have a requirement to find a line in a text file which contains a specific string and then append that line and all lines that follow it to a list. This is the way I accomplished it..

file1 = open("input.txt", "r")
always_print = False
lines = file1.readlines()
output = []
for line in lines:
    if always_print or "def set" in line:  #def set is the string i want
        output.append(line)
        always_print = True

While this works fine,I tried doing the same using list comprehensions.This is what i got :

lines = [ item.strip() for item in open("input.txt")]
always_print = False
output = [item for item in lines if "def set" or print_always in item]

This obviously does not work as I don't set the always_print=True when the desired string is found.How do i do that within the list comprehension ??

Answer 1

Use itertools.dropwhile() to find the first line that contains your def set:

from itertools import dropwhile

output = list(dropwhile(lambda l: 'def set' not in l, lines))

dropwhile() will skip any entry in lines that doesn't match your test; as soon as it matches it stops testing and simply yields everything from there on out.

dropwhile() returns an iterator; I used list() here to convert it to a list of lines, but you could use it as a basis for another loop too, like stripping of newlines, etc.

Demo:

>>> from itertools import dropwhile
>>> lines = '''\
... foo
... bar
... def set():
...     spam
...     ham
...     eggs
... '''.splitlines(True)
>>> list(dropwhile(lambda l: 'def set' not in l, lines))
['def set():\n', '    spam\n', '    ham\n', '    eggs\n']