Supporting containment checks PREMIUM

How can you make your objects work with Python's in operator?

What are containment checks?

Python's in operator is used for containment checks.

Most collections (objects that contain other objects) work with the in operator.

For example, the in operator can be used to check whether a list has a certain element within it:

>>> numbers = [2, 1, 3, 4, 7]
>>> 3 in numbers
True
>>> 5 in numbers
False

All iterables support containment checks

An iterable is an object that can be iterated over.

This Location class makes iterable objects thanks to its __iter__ method:

class Location:
    def __init__(self, lat, long):
        self.lat = lat
        self.long = long

    def __repr__(self):
        return f"Location({self.lat}, {self.long})"

    def __iter__(self):
        yield self.lat
        yield self.long

When we loop over a Location object, we'll see the latitude and the longitude for that location:

>>> melbourne = Location(-37.840935, 144.946457)
>>> list(melbourne)
[-37.840935, 144.946457]

By default, all iterables work with the in operator.

When we ask whether an item is contained in an iterable, Python will loop over that object until it finds the item or reaches the end of the iterable:

>>> -37.840935 in melbourne
True
>>> 100 in melbourne
False

So all iterables support the in operator automatically.

But some of them don't use the default behavior.

For example, when using the in operator on dictionaries, Python doesn't loop all the way through the dictionary looking for a matching key. Instead, it relies on the hashability of dictionary keys to quickly find whether there's a match:

>>> fruit_counts = {"apple": 2, "banana": 4, "mango": 1}
>>> "lime" in fruit_counts
False

Note: sets also support quick containment checks.

And when we use the in operator on strings, Python doesn't loop character by character. Instead, it performs a substring check:

>>> message = "I enjoy writing Python code."
>>> "Python" in message
True

Customizing containment checks on your objects

To customize how the in operator works on our objects, we can implement a __contains__ method.

For example, for our Location class, we may decide that using the in operator should raise an exception because the idea of containment is a little bit unusual with lat-long coordinates:

class Location:
    def __init__(self, lat, long):
        self.lat = lat
        self.long = long

    def __repr__(self):
        return f"Location({self.lat}, {self.long})"

    def __iter__(self):
        yield self.lat
        yield self.long

    def __contains__(self, value):
        raise TypeError("Location objects do not support containment checks")

Note that the in operator and the __contains__ method act a bit differently from other dunder methods.

Most dunder methods, like __eq__ and __add__, are called on the left-hand object for their operator, and the right-hand object is passed into them. The __contains__ method is instead called on the right-hand object, and the left-hand object is passed into it:

>>> melbourne = Location(-37.840935, 144.946457)
>>> melbourne.__contains__(100)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/trey/location.py", line 14, in __contains__
    raise TypeError("Location objects do not support containment checks")
TypeError: Location objects do not support containment checks
>>> 100 in melbourne
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/trey/location.py", line 14, in __contains__
    raise TypeError("Location objects do not support containment checks")
TypeError: Location objects do not support containment checks

That's what the in operator actually does under the hood in Python. It works this way because when we ask whether one object is contained in a collection, we need to ask that question of the collection to find out the answer.

For efficient in lookups, implement __contains__

When we use the in operator between two objects in Python, Python first checks for a __contains__ method on the second object. If it doesn't find a __contains__ method, it falls back to iteration (relying on __iter__ usually).

So, if you're implementing a custom mapping or set in which the in operator shouldn't iterate over your collection, but should instead return True or False quickly, you should implement a __contains__ method.