Itertools Recipe: All Equal

The itertools docs has a ton of slick recipes for using the library to good effect. Some of the code is more useful than illustrative, so I wanted to use these notebooks to break down a few of the functions.

This is

# poor import style, but I want to copy-paste the code
# as-is from the docs

from itertools import *
import itertools

all_equal()

def all_equal(iterable):
    "Returns True if all the elements are equal to each other"
    g = groupby(iterable)
    return next(g, True) and not next(g, False)

Demo

all_equal('aaaaaaaaa')

True

all_equal('aaaaaaab')

False

Why this works

This one relies on how the itertools.groupby() function resolves. Calling it on an iterable yields a stream of tuples of the form (unique group found, iterable).

When they’re all the same, this will be one tuple

obj = itertools.groupby('aaaaaaa')

list(obj)

[('a', <itertools._grouper at 0x25db4853cc0>)]

list(obj)

[('a', <itertools._grouper at 0x1e531b6bf28>)]

and when it gets a different value on the end, it’s got two

obj = itertools.groupby('aaaaaab')

list(obj)

[('a', <itertools._grouper at 0x25db486d240>),
 ('b', <itertools._grouper at 0x25db486d2e8>)]

the default behavior of the next() function is what ties this all together. Calling next() on an iterable with a second argument is very similar to the behavior of dict.get(default_val)

# without default
try:
    obj = iter([])
    print(next(obj))
except StopIteration:
    print('Empty Iter')
    
# with default
try:
    obj = iter([])
    print(next(obj, 0))
except StopIteration:
    print('Empty Iter')
    

Empty Iter
0

So in this context, lets look at that second chunk

and not next(g, False)

If we only have one tuple returned by groupby(), next(g) will raise the StopIterationException, thus giving us the default value False. Applying the leading not evaluates the statement to True.

Otherwise, if it weren’t empty, anything that gets returned will be “Truthy”, get the not applied to it and evaluate to False.