如何正确重置csv.reader以进行多次迭代？

Question

我在使用自定义迭代器时遇到了一个问题，它只能遍历文件一次。我在每次遍历之间都调用了 seek(0) 来重置文件指针，但在第二次遍历时，第一次调用 next() 就抛出了 StopIteration 错误。我觉得我可能忽略了什么明显的东西，但希望能有新的人帮我看看：

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']

Answer 1

DictReader 对象似乎不支持在打开的文件上使用 seek() 命令，所以每次调用 next() 时，都是从文件的末尾开始读取。

在你的 reset 方法中，你可以重新打开文件（你还需要把文件名存储在 self._filename 中）：

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

你也可以考虑像这个问题的最佳答案那样，给你的文件对象创建一个子类。

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

Answer 3

我觉得你最好不要尝试使用 .seek(0)，而是每次都从文件名重新打开文件。

另外，我不建议你在 __iter__() 方法中直接返回 self。这样的话，你的对象就只有一个实例。如果有人在两个不同的线程中使用你的对象，结果可能会很意外。

所以，保存文件名，然后在 __iter__() 方法中，创建一个新的对象，里面有一个新初始化的读取器对象和一个新打开的文件句柄对象；从 __iter__() 返回这个新对象。这样做每次都能正常工作，无论文件对象是什么。它可能是一个从服务器获取数据的网络功能的句柄，或者其他什么，可能不支持 .seek() 方法；但你知道如果重新打开它，你会得到一个新的文件句柄对象。如果有人使用 threading 模块并行运行你的类的10个实例，每个实例都会获取文件中的所有行，而不是随机获取大约十分之一的行。

此外，我不建议你在 MappedIterator 的 .next() 方法中使用异常处理。.__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，这意味着每次调用 .next() 方法时，len(obj) 的值都会改变。如果用户想知道长度，他们应该把结果存储在一个列表中，然后获取列表的长度，或者类似的做法。

编辑：我在 .__iter__() 方法中调用 call_with() 时添加了 **self._kwargs。这样，如果你的 call_with() 函数需要任何额外的参数，它们都会被传递过去。在我做这个更改之前，保存 kwargs 参数在对象中并没有什么好理由；在类的 .__init__() 方法中添加一个 call_with 参数，默认值为 None 也一样。我觉得这个更改是个好主意。