python - 在 python 列表理解中添加额外的语句

Question

我需要在包含特定字符串的文本文件中找到一行，然后将该行及其后面的所有行追加到列表中。这就是我完成它的方式..

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

虽然这工作正常，但我尝试使用列表推导来做同样的事情。这就是我得到的:

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]

这显然不起作用，因为当找到所需的字符串时我没有设置 always_print=True。我如何在列表理解中做到这一点？

Answer 1

使用itertools.dropwhile()找到包含您的 def set 的第一行:

from itertools import dropwhile

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

dropwhile() 将跳过 lines 中与您的测试不匹配的任何条目；一旦匹配，它就会停止测试并从那里简单地产生所有东西。

dropwhile() 返回一个迭代器；我在这里使用 list() 将其转换为行列表，但您也可以将其用作另一个循环的基础，例如去除换行符等。

演示:

>>> 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']