为多次迭代重置 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

我认为您最好不要尝试执行.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__() 方法中将**self._kwargs 添加到对call_with() 的调用中。这样，如果您的 call_with() 函数需要任何额外的参数，它们将被传递。在我进行此更改之前，没有很好的理由将 kwargs 参数保存在对象中；最好将call_with 参数添加到.__init__() 类方法中，默认参数为None。我觉得这个改动很好。

Answer 2

DictReader 对象似乎没有跟随打开文件上的seek() 命令，因此next() 调用是从文件末尾不断进行的。

您可以在reset 中重新打开文件（您还需要将文件名存储在self._filename 中）：

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

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

Answer 3

对于字典阅读器：

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