定义一个支持两个或多个同时迭代的可迭代类

Question

我不确定这个标题好不好。希望有人能帮忙改一下。

我正在尝试定义一个molecule 类，并希望它可以迭代它的原子。 我搜索了人们如何定义一个可迭代的类，它可能看起来像这样：

class molecule(object):
    def __init__(self,name):
         self.__pointer=0
         # ...something else
    def __iter__(self):
         self.__pointer=0
         return self
    def __next__(self):
         self.__pointer+=1
         if self.__pointer<=self.natoms: # self.natoms is defined somewhere
               return self[self.__pointer]   # suppose __getitem__ is defined to return an atom object
         else:
               raise StopIteration("Exceeded all atoms")

效果很好：

 >>> ben=molecule('ben') #defined a molecule
 >>> ben.addatom(...) # and defined some atoms
 >>> ite=iter(ben)
 >>> for atom in ite:
 ...     print(atom)
 ...
 # atom objects are correctly printed here

但是，我发现它不能在两个迭代器中工作，如果它们同时存在的话。

>>> ite=iter(ben)
>>> next(ite)
# atom1
>>> next(ite)
# atom2
>>> ite2=iter(ben)
>>> next(ite)
# atom1 again, where atom3 is expected
>>> next(ite2)
# atom2, where atom1 is expected

这并不奇怪，因为两个迭代器共享同一个self.__pointer，因此定义一个新的迭代器会将指针更新为零。

我查看了此页面How to make class iterable?，其中大多数在课堂上使用self.__pointer，这引发了我的问题。 我猜如果pointer是迭代器的一个属性（ite或ite2）而不是被迭代对象本身（molecule），那么这个问题就可以解决了。

希望有人可以提供帮助:) 谢谢。

Answer 1

使__iter__ 每次被调用时都返回新的迭代器对象。

class molecule_iterator(object):
    def __init__(self, obj):
        self.pointer = 0
        self.obj = obj

    def __next__(self):
        self.pointer += 1  # Are you sure to do this here?
        if self.pointer < self.obj.natoms:
            return self.obj[self.pointer]
        raise StopIteration()


class molecule(object):
    ...

    def __iter__(self):
        return molecule_iterator(self)

    ...

或者，您可以使用yield statement，它不需要您定义__next__。因为__iter__() 会返回一个新的生成器迭代器：

def __iter__(self):
    for i in range(self.natoms):
        yield self[i + 1]

演示：https://ideone.com/TK78Ml