CODEWITHSUNDEEP
pythongenerator
Cribber
Cribber

Reputation: 2913

For-loop first iteration over generator

I have a generator which gives me "chunks" of a pandas dataframe. I save the chunks into a csv file.

For the first "chunk" I want to disable "mode='a'" as I want to overwrite the file if one exists already, the following chunks should be added to the newly created file.

For now I have solved it with a variable outside the loop: "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)

Is there a more elegant way to treat the first element of a generator differently than the rest?

I have found the following code to treat the first object in a list differently, however, it does not work for the generator-object.

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

Trying to implement the code of the list for the generator gives me a TypeError ('generator' object is not subscriptable):

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

Upvotes: 0

Views: 104

Answers (2)

DeepSpace
DeepSpace

Reputation: 81594

I'd use next. This way you don't have the overhead of the repetitive if check (as neglectable as it may be) every iteration:

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)

Upvotes: 2

Tomalak
Tomalak

Reputation: 338128

How about this:

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

Alternatively with a helper variable:

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

Upvotes: 1

Related Questions