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)

生成器上的循环首次迭代

2 个答案: