Первая итерация цикла 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[0].to_csv(filename, sep=';')
for chunk in generator[1:]:
    chunk.to_csv(filename, sep=';', mode='a', header=False)

Answer 1

Как насчет этого:

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)

Answer 2

Я бы использовал 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)