Python Repeatable Generator Decorator

This post goes over creating iterables, generators then using both of these to create a generator decorator that will help you create generators faster.

Iterables

An iterable is any object that can return an iterator using the __iter__ method. An iterator is a stateful object that will produce the next value when you call __next__ on it. This means that any object that contains a __next__ method is an iterator. The important things to remember is that these two objects work together to create an iterative object. Let’s create a simple iterator and iterable object:

class Iterator:
    def __init__(self):
        self.number = 0
    def __next__(self):
        self.number += 1
        if self.a >= 20:
            raise StopIteration

        return self.number

class Iterable:
    def __iter__(self):
        return Iterator()

First we need to create the underlying Iterator that will return a value on __next__. This simple example will increment the number attribute and return the new value. This is where all the work of iteration will be. The Iterable object is just a container to create new Iterator objects. To use this, we just need to create an Iterator, initialize an iterable, then call next on it:

my_iter_class = Iterator()
my_iter = iter(my_iter_class)
print(next(my_iter))
print(next(my_iter))
print(next(my_iter))
print(next(my_iter))

The iter function will take an iterator instance and under the hood call __iter__ to create a new iterable. The next function will then call __next__ on the iterator instance returning a new value each time. You’ll notice that once we reach 20, a StopIteration will be thrown. This is the default error that will tell python that the iteratable is complete. We can combine these two classes to make one object then use it in a simple loop:

class Iterable:
    def __iter__(self):
        self.number = 0
        return self

    def __next__(self):
        self.number += 1
        if self.a >= 20:
            raise StopIteration
        return self.number

    
for x in Iterable():
    print(x)

This will print a number from 0-19 then at 20 throw the StopIteration exception which will be caught in the loop and stop it.

Generators

A generator is a simple function but with a yield state to return instead of return. What makes this function different is that the generator function is paused at the yield and can return to it’s paused state later on. Generators can be treated as an iterator this way:

def counter(num):
    while number < 20:
    yield num
    num += 1

This function is similar to the one in the iterator example except we can pass in an initial counter value and each time this function is called it will return an incremented number until it reaches 20 and ends. When we call counter(0) the function will pause and return a generator object (not the value 0). This generator object can be used just like an iterator with the next() function:

my_iter = counter(0)
print(next(my_iter))
print(next(my_iter))
print(next(my_iter))
print(next(my_iter))

This is the basis for writing generator expressions for lists, dictionaries, and set comprehensions:

my_list = [x*x for x in range(5)]
my_dictionary = {k:v*v for (k,v) in range(5)}
my_set = {x for x in range(5)}

Generator expressions are created the same way but with parenthesis:

my_generator = (x for x in range(5))

Generator Decorator

We can combine these two concepts with decorators I discussed in a previous post. We can take a generator function as the input to a decorator and output an iterable object:

def repeatable(generator):
    class RepeatableGenerator:
        def __init__(self, *args, **kwargs):
            self.args = args
            self.kwargs = kwargs
        def __iter__(self):
            return iter(generator(*self.args, **self.kwargs))
    return RepeatableGenerator

From here we can use it on generator functions:

@repeatable
def generator(max):
    for x in range(max):
        yield x

g = generator(5)
list(g) # [0, 1, 2, 3, 4]

h = generator(5)
list(h) # [0, 1, 2, 3, 4]

Each time you call the generator function it will return a new iterable so you don’t have to worry about states impacting each other. You’ll notice a small difference here though to our example decorators in my previous post: this decorator returns a class instead of a function. This can causes problems with code that expect the returned object to be a function instead of a class. This can include when we try to wrap multiple decorators with this decorator – function decorators expect a function. The class returned also can’t be used as a method to another class since it doesn’t have a __get__ method to bind to the owner class of instance of it. To fix this we can make a wrapper function that will return a new instance of the class instead. functools.wraps is there to keep internal function attributes the same (also discussed in my previous post of decorators).

import functools
def repeatable(generator):

    class RepeatableGenerator:
        def __init__(self, *args, **kwargs):
            self.args = args
            self.kwargs = kwargs
        def __iter__(self):
            return iter(generator(*self.args, **self.kw

    @functools.wraps(generator)
    def wrapper(*args, **kwargs):
        return RepeatableGenerator(*args, **kwargs)
    return wrapper