CODEWITHSUNDEEP
pythoncsvgeneratorwriter
uncrase
uncrase

Reputation: 694

Don't write file with csv DictWriter if there is no data

I have a Python (3.4) routine that writes a csv file using a generator. However, depending on the parameter set, there may not be any data, in which case, I don't want the csv file to be written. (It would just write the file with a header only).

Right now, the bandaid is to count the lines after generation and then delete the file, but surely there must be a better way, while retaining the pattern of having a generator being the only code that's aware of whether there is data for the given parameters, (nor having to call on the generator twice):

def write_csv(csv_filename, fieldnames, generator, from_date, to_date, client=None):
  with open(csv_filename, 'w', newline='') as csv_file:
    csv_writer = csv.DictWriter(csv_file, fieldnames=fieldnames, delimiter='\t')
    csv_writer.writeheader()
    csv_writer.writerows(generator(from_date, to_date, client))

  # If no rows were written delete the file, we don't want it
  with open(csv_filename) as f:
    lines = sum(1 for _ in f)
    if lines == 1:
      f.close()
      os.remove(f.name)


def per_client_items_generator(from_date, to_date, client):
  return (per_client_detail(client, sales_item) for sales_item in
        sales_by_client.get(client))

Upvotes: 2

Views: 703

Answers (2)

Matthias Fripp
Matthias Fripp

Reputation: 18625

You could use itertools to take a look at the first item and then sort of put it back into the generator:

import itertools
gen = generator(from_date, to_date, client)
try:
    # try to get an element
    first = next(gen)
except StopIteration:
    pass
else:
    # run this if there was no exception:
    gen = itertools.chain([first], gen)
    csv_writer.writeheader()
    csv_writer.writerows(gen)

This is a little shorter, but may be harder to read:

import itertools
gen = generator(from_date, to_date, client)
try:
    # pop an element then chain it back in
    gen = itertools.chain([next(gen)], gen)
except StopIteration:
    pass
else:
    # run this if there was no exception:
    csv_writer.writeheader()
    csv_writer.writerows(gen)

Or this doesn't use visible try/catch code (although there's probably an equal amount down inside next()):

import itertools
sentinel = object()  # special flag that couldn't come from the generator
gen = generator(from_date, to_date, client)

# try to get something
first = next(gen, sentinel)
if first is not sentinel:
    # got a meaningful item, put it back in the generator
    gen = itertools.chain([first], gen)
    csv_writer.writeheader()
    csv_writer.writerows(gen)

(These were inspired by Stephen Rauch's answer, but with a few tweaks.)

Upvotes: 2

Stephen Rauch
Stephen Rauch

Reputation: 49814

You can sort of preview the generator using next(), being careful to preserve the first generated value, with something like:

csv_gen = generator(from_date, to_date, client)
try:
    first_item = next(csv_gen)
except StopIteration:
    csv_gen = None

if csv_gen is not None:
     # prep for write csv
     ....         

    # write csv header
    csv_writer.writeheader()

    # write item already read from generator
    csv_writer.writerow(first_item)        

    # write rest of generator
    csv_writer.writerows(csv_gen)

Please note that this was not tested, so may contain silly typos.

Upvotes: 1

Related Questions