生成器上的 for 循环第一次迭代

Question

我有一个生成器，它为我提供了熊猫数据框的“块”。我将块保存到 csv 文件中。

对于第一个“块”，我想禁用“mode='a'”，因为如果文件已经存在，我想覆盖该文件，应将以下块添加到新创建的文件中。

现在我已经用循环外的一个变量解决了它：“first”。

first = True
for chunk in generator:
    if first:
        chunk.to_csv(filename, sep=';')
        first = False
    else:
        chunk.to_csv(filename, sep=';', mode='a', header=False)

有没有更优雅的方式来处理生成器的第一个元素而不是其他元素？

---

我发现以下代码可以不同地处理列表中的第一个对象，但是它不适用于生成器对象。

seq= something.get()
foob( seq[0] )
for member in seq[1:]:
    foo( member )

尝试为生成器实现列表的代码给了我一个 TypeError（'generator' 对象不可下标）：

generator[0].to_csv(filename, sep=';')
for chunk in generator[1:]:
    chunk.to_csv(filename, sep=';', mode='a', header=False)

Answer 1

我会使用next。这样一来，您就不必在每次迭代时重复进行 if 检查（尽管可以忽略）：

first = next(generator)
first.to_csv(filename, sep=';')
for chunk in generator:  # will start from the second element
    chunk.to_csv(filename, sep=';', mode='a', header=False)

Answer 2

这个怎么样：

for i, chunk in enumerate(generator):
    chunk.to_csv(filename, sep=';', mode=('w' if i == 0 else 'a'), header=(i == 0))

或者使用辅助变量：

for i, chunk in enumerate(generator):
    first = i == 0
    chunk.to_csv(filename, sep=';', mode=('w' if first else 'a'), header=first)