🔥 Iterator

Machine learning uses iterators extensively (e.g. for Dataloaders). So lets explore how iterators work in Mojo and build one ourself.

Firstly let’s look at how to implement a simple iterator in Python and then use that as a starting point for a Mojo implementation.

class Counter:
    def __init__(self, low, high):
        self.current = low - 1
        self.high = high

    def __iter__(self):
        return self

    def __next__(self): # Python 2: def next(self)
        self.current += 1
        if self.current < self.high:
            return self.current
        raise StopIteration


for c in Counter(3, 9):
    print(c)

3
4
5
6
7
8

Now in Mojo it is similar to Python in that we define a __iter__ and a __next__. However, Mojo doesn’t understand the StopIteration exception (that the __next__ in the Python implementation returns)! So how do we tell Mojo that the Iterator has reached the end? As per the Mojo Changelog, the control flow exits automatically when the length is zero.
So, we need to implement a __len__ and make sure it decrements with every call to __next__.

@value
struct Counter:
    var current: Int
    var len: Int

    fn __init__(inout self, low: Int, high: Int):
        self.current = low - 1
        self.len = high - low

    fn __len__(self) -> Int:
        return self.len

    fn __iter__(self) -> Self:
        return self

    fn __next__(inout self) -> Int: 
        self.len = self.len - 1
        self.current = self.current + 1
        return self.current
    
for c in Counter(3, 9):
    print(c)

3
4
5
6
7
8

Note that the @value decorator is used to avoid the boilerplate for the __copyinit__ and __moveinit__ methods. See mojodojo.dev for a short tutorial on the @value decorator.