2

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))
Stephen Rauch
  • 47,830
  • 31
  • 106
  • 135
uncrase
  • 644
  • 1
  • 8
  • 16
  • You might also take a look at http://stackoverflow.com/questions/661603/how-do-i-know-if-a-generator-is-empty-from-the-start or http://stackoverflow.com/questions/3114252/one-liner-to-check-whether-an-iterator-yields-at-least-one-element. – Matthias Fripp May 08 '17 at 22:32

2 Answers2

2

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.)

Matthias Fripp
  • 17,670
  • 5
  • 28
  • 45
  • Yeah, but now I'm worried about how to avoid catching StopIteration errors that bubble up somehow from the code after first(). Might have to define a helper variable... – Matthias Fripp May 08 '17 at 22:12
  • I think that's basically what I did. For a minute I thought I'd need to set a helper variable after getting first(), to indicate whether it succeeded. Then I remembered `try`/`except`/`else`, which is perfect for this. – Matthias Fripp May 08 '17 at 22:16
  • I was originally going with the else in my code, but wasn't sure where the OP was at, so went a bit simpler. The Chain is clever... – Stephen Rauch May 08 '17 at 22:17
  • 1
    Thanks! Just noticed that a couple of the answers at http://stackoverflow.com/questions/661603/how-do-i-know-if-a-generator-is-empty-from-the-start are similar to mine, but so it goes! – Matthias Fripp May 08 '17 at 22:33
1

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.

Stephen Rauch
  • 47,830
  • 31
  • 106
  • 135