Membership: __contains__, __iter__, and __getitem__

The iteration story is even richer than we’ve seen thus far. Operator overloading is often layered: classes may provide specific methods, or more general alternatives used as fallback options. For example:

In the iterations domain, classes can implement the in membership operator as an iteration, using either the __iter__ or __getitem__ methods. To support more specific membership, though, classes may code a __contains__ method—when present, this method is preferred over __iter__, which is preferred over __getitem__. The __contains__ method should define membership as applying to keys for a mapping (and can use quick lookups), and as a search for sequences.

Consider the following class, whose file has been instrumented for dual 2.X/3.X usage using the techniques described earlier. It codes all three methods and tests membership and various iteration contexts applied to an instance. Its methods print trace messages when called:

# File contains.py
from __future__ import print_function         # 2.X/3.X compatibility

class Iters:
    def __init__(self, value):
        self.data = value
    
    def __getitem__(self, i):                 # Fallback for iteration
        print('get[%s]:' % i, end='')         # Also for index, slice
        return self.data[i]
    
    def __iter__(self):                       # Preferred for iteration
        print('iter=> ', end='')              # Allows only one active iterator
        self.ix = 0
        return self
    
    def __next__(self):
        print('next:', end='')
        if self.ix == len(self.data): raise StopIteration
        item = self.data[self.ix]
        self.ix += 1
        return item
    
    def __contains__(self, x):                # Preferred for 'in'
        print('contains: ', end='')
        return x in self.data
    next = __next__                           # 2.X/3.X compatibility

if __name__ == '__main__':
    X = Iters([1, 2, 3, 4, 5])        # Make instance
    print(3 in X)                     # Membership
    for i in X:                       # for loops
        print(i, end=' | ')

    print()
    print([i ** 2 for i in X])        # Other iteration contexts
    print( list(map(bin, X)) )

    I = iter(X)                       # Manual iteration (what other contexts do)
    while True:
        try:
            print(next(I), end=' @ ')
        except StopIteration:
            break

As is, the class in this file has an __iter__ that supports multiple scans, but only a single scan can be active at any point in time (e.g., nested loops won’t work), because each iteration attempt resets the scan cursor to the front. Now that you know about yield in iteration methods, you should be able to tell that the following is equivalent but allows multiple active scans—and judge for yourself whether its more implicit nature is worth the nested-scan support and six lines shaved (this is in file contains_yield.py):

class Iters:
    def __init__(self, value):
        self.data = value
    
    def __getitem__(self, i):                 # Fallback for iteration
        print('get[%s]:' % i, end='')         # Also for index, slice
        return self.data[i]
    
    def __iter__(self):                       # Preferred for iteration
        print('iter=> next:', end='')         # Allows multiple active iterators
        for x in self.data:                   # no __next__ to alias to next
            yield x
            print('next:', end='')
    
    def __contains__(self, x):                # Preferred for 'in'
        print('contains: ', end='')
        return x in self.data

On both Python 3.X and 2.X, when either version of this file runs its output is as follows—the specific __contains__ intercepts membership, the general __iter__ catches other iteration contexts such that __next__ (whether explicitly coded or implied by yield) is called repeatedly, and __getitem__ is never called:

contains: True
iter=> next:1 | next:2 | next:3 | next:4 | next:5 | next:
iter=> next:next:next:next:next:next:[1, 4, 9, 16, 25]
iter=> next:next:next:next:next:next:['0b1', '0b10', '0b11', '0b100', '0b101']
iter=> next:1 @ next:2 @ next:3 @ next:4 @ next:5 @ next:

Watch what happens to this code’s output if we comment out its __contains__ method, though—membership is now routed to the general __iter__ instead:

iter=> next:next:next:True
iter=> next:1 | next:2 | next:3 | next:4 | next:5 | next:
iter=> next:next:next:next:next:next:[1, 4, 9, 16, 25]
iter=> next:next:next:next:next:next:['0b1', '0b10', '0b11', '0b100', '0b101']
iter=> next:1 @ next:2 @ next:3 @ next:4 @ next:5 @ next:

And finally, here is the output if both __contains__ and __iter__ are commented out—the indexing __getitem__ fallback is called with successively higher indexes until it raises IndexError, for membership and other iteration contexts:

get[0]:get[1]:get[2]:True
get[0]:1 | get[1]:2 | get[2]:3 | get[3]:4 | get[4]:5 | get[5]:
get[0]:get[1]:get[2]:get[3]:get[4]:get[5]:[1, 4, 9, 16, 25]
get[0]:get[1]:get[2]:get[3]:get[4]:get[5]:['0b1', '0b10', '0b11', '0b100','0b101']
get[0]:1 @ get[1]:2 @ get[2]:3 @ get[3]:4 @ get[4]:5 @ get[5]:

As we’ve seen, the __getitem__ method is even more general: besides iterations, it also intercepts explicit indexing as well as slicing. Slice expressions trigger __getitem__ with a slice object containing bounds, both for built-in types and user-defined classes, so slicing is automatic in our class:

>>> from contains import Iters
>>> X = Iters('spam')               # Indexing
>>> X[0]                            # __getitem__(0)
get[0]:'s'

>>> 'spam'[1:]                      # Slice syntax
'pam'
>>> 'spam'[slice(1, None)]          # Slice object
'pam'

>>> X[1:]                           # __getitem__(slice(..))
get[slice(1, None, None)]:'pam'
>>> X[:-1]
get[slice(None, −1, None)]:'spa'

>>> list(X)                         # And iteration too!
iter=> next:next:next:next:next:['s', 'p', 'a', 'm']

In more realistic iteration use cases that are not sequence-oriented, though, the __iter__ method may be easier to write since it must not manage an integer index, and __contains__ allows for membership optimization as a special case.