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

2024-03-29 02:38:06 发布

您现在位置:Python中文网/ 问答频道 /正文

自定义迭代器有一个问题,因为它只在文件上迭代一次。我在两次迭代之间对相关文件对象调用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']

Tags: selfobjforreturndefargsvalmapping
3条回答

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

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

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

另外,我不建议在MappedIterator中的.next()方法内使用异常处理程序。.__iter__()方法应该返回一个可以可靠迭代的对象。如果一个愚蠢的用户传入一个整数对象(例如:3),这将是不可iterable的。在.__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()的值都将更改。如果用户想知道长度,他们应该将结果存储在一个列表中,并获取列表的长度,或者类似的内容。

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

对象DictReader似乎不遵循打开文件上的seek()命令,因此next()调用从文件末尾开始连续进行。

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

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

您还可以将文件对象的子类化类似于this问题的顶级答案。

对于听写器:

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

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

对于听写器:

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

相关问题 更多 >