使用生成器读取CSV文件

时间:2018-11-10 02:18:29

标签: python python-3.x csv iterator generator

我有一个很大的csv文件需要处理,并且这种方式已经完成(非常简化):

import csv
from csv import excel

def _get_dialect():
    class CustomDialect(excel):
        delimiter = ','
    return CustomDialect()

class DictIter:
    def __init__(self):
        self.reader = csv.DictReader(open('test.csv'), 
                                     fieldnames=['col1', 'col2'], 
                                     dialect=_get_dialect())

    def __iter__(self):
        return self

    def __next__(self):
        return self.reader.__next__()

items = DictIter()
for item in items:
    print(item)

有更好的方法可以做到这一点,但这就是我现在所拥有的。 现在,我想将处理分成多个部分,并找到适合我的简单解决方案:

def gen_chunks(reader, chunksize=500):
    chunk = []
    for i, line in enumerate(reader):
        if (i % chunksize == 0 and i > 0):
            yield chunk
            del chunk[:]
        chunk.append(line)
    yield chunk

但是,我缺少一些生成器知识来将这两段代码结合在一起。基本上我想要这样的东西:

import csv
from csv import excel

def gen_chunks(reader, chunksize=500):
    chunk = []
    for i, line in enumerate(reader):
        if (i % chunksize == 0 and i > 0):
            yield chunk
            del chunk[:]
        chunk.append(line)
    yield chunk

def _get_dialect():
    class CustomDialect(excel):
        delimiter = ','
    return CustomDialect()

class DictIter:
    def __init__(self):
        self.reader = csv.DictReader(open('test.csv'), 
                                     fieldnames=['col1', 'col2'], 
                                     dialect=_get_dialect())

    def __iter__(self):
        return self

    def __next__(self):
        for chunk in gen_chunks(self.reader):
            for item in chunk:
                yield item

items = DictIter()
for item in items:
    print(item)

这可能有点笨拙,但是我想在对当前结构进行最小更改的情况下实现拆分处理。我要实现的目标是让当前实现保留迭代器类,但一次处理一个块,并在处理完上一个块后产生下一个块。

1 个答案:

答案 0 :(得分:1)

您的解决方案看起来过于复杂:

import csv

with open('test.csv',newline='') as f:
    reader = csv.DictReader(f,fieldnames=['col1', 'col2'])
    for item in reader:
        print(item)

您的分块仍然一次返回一项。 excel方言的默认定界符是逗号,excel是默认方言。

注意newline=''是打开传递给csv读取器或写入器对象的文件的记录方法。

如果您有真正的理由进行分块(多处理?),则应声明并在尝试不起作用时显示尝试。