pythonpython-3.xclassfor-loop

How to make a custom class work with the ‘for in’ loop in Python?


If I have a custom class, how can I make it work with a ‘for in’ loop, like this:

class number_range:
    def __init__(self, start, end):
        self.__data = []
        for i in range(start, end):
            self.__data.append(i)
    #something that will allow the for in loop

numbers = number_range(1, 100)          
for number in numbers:
    print(number)

Solution

  • What you are looking for is an iterator. For example, this one acts like range but only gives you odd numbers:

    class oddrange(object):
        def __init__(self, first, lastPlusOne):
            # Ensure start at first odd >= given start, store terminator.
    
            self.__curr = first + 1 - first % 2
            self.__term = lastPlusOne
    
        def __iter__(self):
            return self
    
        def __next__(self):
            # Terminate if range over, otherwise return current, calculate next.
    
            if self.__curr >= self.__term:
                raise StopIteration()
    
            (cur, self.__curr) = (self.__curr, self.__curr + 2)
            return cur
    

    You can see this in action with the calling code:

    for i in oddrange(-5, 9): print(i, end = ' ')
    print()
    for i in oddrange(-3, 9): print(i, end = ' ')
    print()
    for i in oddrange(-4, 10): print(i, end = ' ')
    print()
    

    This generates:

    -5 -3 -1 1 3 5 7
    -3 -1 1 3 5 7
    -3 -1 1 3 5 7 9