如果没有数据，请勿使用 csv DictWriter 写入文件

Question

我有一个 Python (3.4) 例程，它使用生成器编写 csv 文件。但是，根据参数集，可能没有任何数据，在这种情况下，我不希望写入csv 文件。 （它只会写入带有标题的文件）。

目前，创可贴是在生成后计算行数，然后删除文件，但肯定有更好的方法，同时保留生成器是唯一知道是否有数据的代码的模式给定的参数，（也不必调用生成器两次）：

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

Answer 1

您可以使用 itertools 查看第一项，然后将其放回生成器中：

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)

这有点短，但可能更难阅读：

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)

或者这不使用可见的try/catch 代码（尽管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)

（这些灵感来自 Stephen Rauch 的回答，但做了一些调整。）

Answer 2

您可以使用next() 对生成器进行预览，小心保留第一个生成的值，例如：

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)

请注意，这未经测试，因此可能包含愚蠢的错别字。