Example python implementation of generator/yield

Question

Let's take the following python code to generate an infinite counter:

def infinite_sequence():
    num = 0
    while True:
        yield num
        num += 1

I can then call this as:

for i in infinite_sequence():
    print (i)

And do a KeyboardInterrupt whenever I want to exit it. The yield/generator stuff is implemented in C/Cpython. What would be an example of implementing a similar "generator functionality" within python? That is, if I could call a function/method in place of the yield num line, and get the same functionality, what would be an example implementation of this in python?

Answer 1

You would have to implement the iterator protocol yourself.

class InfiniteSequence:
    def __init__(self):
        self.x = 0

    def __iter__(self):
        return self

    def __next__(self):
        x = self.x
        self.x = x + 1
        return x


for i in InfiniteSequence():
    print(i)

The __iter__ method has to return something that implements __next__. The __next__ method has to return the next value in the sequence. The instance attribute maintains state between calls to __next__.

The same class doesn't have to implement both methods. You might want a separate iterator class so that you can have multiple, independent iterators over the iterable sequence.

class InfiniteSequence:
    def __init__(self, start):
        self.start = start

    def __iter__(self):
        return InfSeqIterator(self.start)


class InfSeqIterator:
    def __init__(self, x):
        self.x = x

    def __iter__(self):
        return self

    def __next__(self):
        x = self.x
        self.x += 1
        return x


nats = InfiniteSequence()

i1 = iter(nats)
i2 = iter(nats)

assert next(i1) == 0
assert next(i2) == 0  # not 1

---

generator is just one example of a class that implements this protocol; you create an instance of generator using either a generator expression ((x for x in [1,2,3])) or a generator function (i.e., a function that uses yield).