如果没有数据，不要用 csv DictWriter 写入文件

Question

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

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

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

您可以使用 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)

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

Answer 2

您可以使用 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 的回答，但有一些调整。）

如果没有数据，不要用 csv DictWriter 写入文件

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

python

csv

generator

writer