Iteators for Free

21 August 2024

Suppose you are creating a custom collection type and would like to be able to traverse it with a for loop. All you need to do is define the magic method __getitem__. We will show an example here by creating an class to represent a deck of cards. First, we create the Card module. Nothing fancy going on here. Note that we create static variables to store the ranks and suits for cards.


class Card:
    ranks = [str(k) for k in range(2,11)] + ["Jack", "Queen", "King", "Ace"]
    suits = ["clubs", "diamonds", "hearts", "spades"]
    def __init__(self, n):
        self.n = n
    
    def __str__(self):
        return f"{Card.ranks[self.n%13]} of {Card.suits[self.n//13]}"

    def __repr__(self):
        return f"Card({self.n})"

    @property
    def rank(self):
        return Card.ranks[self.n%13]
        
    @property
    def suit(self):
        return Card.ranks[self.n//13]

if __name__ == "__main__":
    print(Card(5))
        

Now we create the Deck module with __getitem__.


import random
from Card import Card
class Deck:
    def __init__(self):
        self.cards = [Card(n) for n in range(52)]
        random.shuffle(self.cards)
    
    def __getitem__(self, index:int):
        return self.cards[index]

    def __len__(self):
        return len(self.cards)

def main():
    d = Deck()
    print(len(d))
    for k in d[:5]:
        print(k)
    print("*"*50)
    for k in d:
        print(k)
if __name__ == "__main__":
    main()
        
        

You can run this and see that you can slice a Deck and print it with a for loop. Now let us inspect it in an interactive Python shell.

>>> from Deck import Deck
>>> d = Deck()
>>> iter(d)
<iterator object at 0x104377430>
        

We get an iterator for free by just implementing __getitem__!

This used to be the only way to make a collection traversable by a for loop back in the Fred 'n Barney primitive Python 2.1 days. Later, the __iter__ method was added. Behind the scenes, this is what happens.

  1. Python looks for an __iter__ method. If it finds one, the call to the iter function is made using the method.
  2. If there is no __iter__, Python calls __getitem__ starting at 0, proceeding an index at a time, and ends when an IndexError is emitted. In this way, it controls a for loop.

A Laginappe Also notice this nice feature from the use of the @property decorator.

        >>> c = Card(25)
        >>> c.rank
        'Ace'
        >>> c.rank = "quack"
        Traceback (most recent call last):
          File "<stdin>", line 1, in <module>
        

Python hisses when a foolish end-user of Our Pristine Code indulges in vandalistic stupidity.