正确的方法来重置csv.reader多次迭代?

时间:2011-07-19 23:46:35

标签: python csv iterator

自定义迭代器存在问题,因为它只会迭代文件一次。我在迭代之间对相关文件对象调用seek(0),但在第二次运行时第一次调用StopIteration时会抛出next()。我觉得我忽略了一些显而易见的事情,但我会对此有一些新的看法:

class MappedIterator(object):
    """
    Given an iterator of dicts or objects and a attribute mapping dict, 
    will make the objects accessible via the desired interface.

    Currently it will only produce dictionaries with string values. Can be 
    made to support actual objects later on. Somehow... :D
    """

    def __init__(self, obj=None, mapping={}, *args, **kwargs):

        self._obj = obj
        self._mapping = mapping
        self.cnt = 0

    def __iter__(self):

        return self

    def reset(self):

        self.cnt = 0

    def next(self):

        try:

            try:
                item = self._obj.next()
            except AttributeError:
                item = self._obj[self.cnt]

            # If no mapping is provided, an empty object will be returned.
            mapped_obj = {}

            for mapped_attr in self._mapping:

                attr = mapped_attr.attribute
                new_attr = mapped_attr.mapped_name

                val = item.get(attr, '')
                val = str(val).strip() # get rid of whitespace

                # TODO: apply transformers...

                # This allows multi attribute mapping or grouping of multiple
                # attributes in to one.
                try:
                    mapped_obj[new_attr] += val
                except KeyError:
                    mapped_obj[new_attr] = val

            self.cnt += 1

            return mapped_obj

        except (IndexError, StopIteration):

            self.reset()
            raise StopIteration


class CSVMapper(MappedIterator):

    def __init__(self, reader, mapping={}, *args, **kwargs):

        self._reader = reader
        self._mapping = mapping

        self._file = kwargs.pop('file')

        super(CSVMapper, self).__init__(self._reader, self._mapping, *args, **kwargs)

    @classmethod
    def from_csv(cls, file, mapping, *args, **kwargs):

        # TODO: Parse kwargs for various DictReader kwargs.
        return cls(reader=DictReader(file), mapping=mapping, file=file)

    def __len__(self):

      return int(self._reader.line_num)

    def reset(self):

      if self._file:

        self._file.seek(0)

      super(CSVMapper, self).reset()

样本用法:

file = open('somefile.csv', 'rb') # say this file has 2 rows + a header row

mapping = MyMappingClass() # this isn't really relevant

reader = CSVMapper.from_csv(file, mapping)

# > 'John'
# > 'Bob'
for r in reader:

  print r['name']

# This won't print anything
for r in reader:

  print r['name']

3 个答案:

答案 0 :(得分:9)

我认为你最好不要尝试.seek(0),而是每次都从文件名中打开文件。

我建议您不要在self方法中返回__iter__()。这意味着您只有一个对象实例。我不知道有人试图从两个不同的线程中使用你的对象的可能性有多大,但如果发生了这种情况,那么结果就会令人惊讶。

因此,保存文件名,然后在__iter__()方法中,使用刚刚初始化的读者对象和新打开的文件句柄对象创建一个新对象;从__iter__()返回此新对象。无论文件类对象到底是什么,这都会每次都有效。它可以是从服务器提取数据或谁知道什么的网络功能的句柄,它可能不支持.seek()方法;但是你知道,如果你再次打开它,你将获得一个新的文件句柄对象。如果有人使用threading模块并行运行您的类的10个实例,则每个实例将始终从文件中获取所有行,而不是每个随机获得大约十分之一的行。

此外,我不建议.next()MappedIterator方法中的异常处理程序。 .__iter__()方法应该返回一个可以可靠迭代的对象。如果一个愚蠢的用户传入一个整数对象(例如:3),那么这将是不可迭代的。在.__iter__()内,你总是可以在一个参数上显式调用iter(),如果它已经是一个迭代器(例如,一个打开的文件句柄对象),你将获得相同的对象;但如果它是一个序列对象,你将获得一个对序列起作用的迭代器。现在,如果用户传入3,则对iter()的调用将引发一个异常,该异常在用户传递3的行处有意义,而不是来自第一次调用.next()的异常。作为奖励,您不再需要cnt成员变量,并且您的代码会更快一些。

所以,如果你把我的所有建议放在一起,你可能会得到这样的结论:

class CSVMapper(object):
    def __init__(self, reader, fname, mapping={}, **kwargs):
        self._reader = reader
        self._fname = fname
        self._mapping = mapping
        self._kwargs = kwargs
        self.line_num = 0

    def __iter__(self):
        cls = type(self)
        obj = cls(self._reader, self._fname, self._mapping, **self._kwargs)
        if "open_with" in self._kwargs:
            open_with = self._kwargs["open_with"]
            f = open_with(self._fname, **self._kwargs)
        else:
            f = open(self._fname, "rt")
        # "itr" is my standard abbreviation for an iterator instance
        obj.itr = obj._reader(f)
        return obj

    def next(self):
        item = self.itr.next()
        self.line_num += 1

        # If no mapping is provided, item is returned unchanged.
        if not self._mapping:
            return item  # csv.reader() returns a list of string values

        # we have a mapping so make a mapped object
        mapped_obj = {}

        key, value = item
        if key in self._mapping:
            return [self._mapping[key], value]
        else:
            return item

if __name__ == "__main__":
    lst_csv = [
        "foo, 0",
        "one, 1",
        "two, 2",
        "three, 3",
    ]

    import csv
    mapping = {"foo": "bar"}
    m = CSVMapper(csv.reader, lst_csv, mapping, open_with=iter)

    for item in m: # will print every item
        print item

    for item in m: # will print every item again
        print item

现在.__iter__()方法每次调用时都会为您提供一个新对象。

请注意示例代码如何使用字符串列表而不是打开文件。在此示例中,您需要指定要使用的open_with()函数,而不是默认的open()来打开文件。由于我们可以迭代我们的字符串列表以一次返回一个字符串,因此我们可以在这里使用iter作为open_with函数。

我不明白你的映射代码。 csv.reader返回一个字符串值列表,而不是某种字典,所以我写了一些简单的映射代码,适用于有两列的CSV文件,第一列是字符串。显然,你应该删除我的普通映射代码并输入所需的映射代码。

另外,我拿出了.__len__()方法。当你执行len(obj)之类的操作时,它会返回序列的长度;你让它返回line_num,这意味着每次调用len(obj)方法时.next()的值都会改变。如果用户想要知道长度,他们应该将结果存储在列表中并获取列表的长度,或类似的东西。

编辑:我在**self._kwargs方法中对call_with()的调用添加了.__iter__()。这样,如果您的call_with()函数需要任何额外的参数,它们将被传递。在我做这个改变之前,没有一个很好的理由在对象中保存kwargs参数;将call_with参数添加到类.__init__()方法,默认参数为None也同样出色。我认为这种改变很好。

答案 1 :(得分:2)

对于DictReader:

f = open(filename, "rb")
d = csv.DictReader(f, delimiter=",")

f.seek(0)
d.__init__(f, delimiter=",")

对于DictWriter:

f = open(filename, "rb+")
d = csv.DictWriter(f, fieldnames=fields, delimiter=",")

f.seek(0)
f.truncate(0)
d.__init__(f, fieldnames=fields, delimiter=",")
d.writeheader()
f.flush()

答案 2 :(得分:1)

DictReader对象似乎没有在打开的文件上遵循seek()命令,因此next()调用将不断从文件末尾进行。

在您的reset中,您可以重新打开该文件(您还需要将文件名存储在self._filename中):

def reset(self):
     if self._file:
         self._file.close()
         self._file = open(self._filename, 'rb')

您还可以查看对文件对象进行子类化,类似于this问题的最佳答案。