MacSanhe
MacSanhe

Reputation: 2240

Python: Advance __iter__ behavior

I have a size-limited dictionary class, and I want to make the iter method works like this:

if the value not None, then generate. Otherwise, skip it. Do you know how to implement it guys?

class myclass(object):
    def __init__(self):
        self.data = {1:'a', 2:None, 3:'c'}

    def __iter__(self):
        return iter(self.data.values())

    def __next__(self): ## <== I THINK THIS IS A WRONG EXAMPLE
        if iter(self): ## BUT I DON"T KNOW HOW TO FIX IT
            return iter(self)

mc = myclass()
for i in mc:
    print(i)

Upvotes: 0

Views: 105

Answers (2)

Ethan Furman
Ethan Furman

Reputation: 69288

To create a single-use iterator:

class myclass(object):
    def __init__(self):
        self.data = {1:'a', 2:None, 3:'c'}
        self.keys = self.data.keys()
        self.index = 0
    def __iter__(self):
        return self
    def __next__(self):
        index = self.index
        while 'searching for value':
            if index >= len(self.keys):
                raise StopIteration
            key = self.keys[index]
            value = self.data[key]
            index += 1
            if value is not None:
                self.index = index
                return value

As you can see, using Martjin's answer is much nicer.

Upvotes: 0

Martijn Pieters
Martijn Pieters

Reputation: 1125078

If your __iter__ method directly returns an iterator, you do not need to implement __next__; it will not be consulted in that case (it is the __next__ method of the returned iterator that is used instead).

Return a generator expression:

class myclass(object):
    def __init__(self):
        self.data = {1:'a', 2:None, 3:'c'}

    def __iter__(self):
        return (v for v in self.data.values() if v is not None)

Demo:

>>> class myclass(object):
...     def __init__(self):
...         self.data = {1:'a', 2:None, 3:'c'}
...     def __iter__(self):
...         return (v for v in self.data.values() if v is not None)
... 
>>> list(myclass())
['a', 'c']

You'd have to return self from __iter__ if you wanted the class to be its own iterator; this would mean you'd need to track state between __next__ calls to know what value to return from each call. For your usecase, you most probably do not want that.

Upvotes: 3

Related Questions