4

I'm not sure if this title is good..Hope someone could help to change it.

I'm trying to define an molecule class, and hope it can iterate over its atoms. I searched how people define an iterable class, and it may look like this:

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

It works well:

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

However, I found it cannot work in two iterators, if they exist simultaneously.

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

This is not a surprise because the two iterators share a same self.__pointer, so defining a new iterator will renew the pointer to zero.

I had a look at this page How to make class iterable?, and most of them use a self.__pointer within the class, which raises my question. I guess if the pointer be an attribute of the iterator(ite or ite2) but not of the iterated object itself (molecule), then this problem could be solved.

Hope someone can give some help:) Thanks.

Ruixing Wang
  • 75
  • 1
  • 7

1 Answers1

2

Make __iter__ to return new iterator object every time when it's called.

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)

    ...

Or, you can use yield statement which does not require you to define __next__. Because __iter__() will return a new generator iterator:

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

DEMO: https://ideone.com/TK78Ml

falsetru
  • 357,413
  • 63
  • 732
  • 636
  • > Make `__iter__` to return new iterator object every time when it's called. << How to do this? while using `yield` looks good. – Ruixing Wang Apr 01 '16 at 07:49
  • @RuixingWang, I updated the answer to include the code for the first option. (return new iterator instance) – falsetru Apr 01 '16 at 07:58
  • Cool answer. For the comment `#Are you sure`, yes, since I count atoms from 1 but not zero:) Thanks! – Ruixing Wang Apr 01 '16 at 08:06