The most powerful way to enter loops in Python

When debugging or just writing exploratory code, you often want to step through the lines of your loop code. If your iterable is a list, that is simple, as you can use indexing to get to the first item:

fruits = ['apple', 'pear', 'plum']
fruit = fruits[0]

for fruit in fruits:
    print(fruit)

But for dictionaries that won’t work. Same for generators.

fruit_dict = dict(apple='red', pear='green', plum='purple')
fruit = fruit_dict[0]     # KeyError: 0

fruit_generator = (fruit for fruit in fruits)
fruit = fruit_generator[0]  # TypeError

You can even mess up your original variables if you were using a defaultdict, as it will add a “0” element to the dictionary.

from collections import defaultdict
fruit_defaultdict = defaultdict(str, dict(apple='red', pear='green', plum='purple'))
fruit = fruit_defaultdict[0]        # ''

You can also fool yourself into thinking you solved the problem by figuring out one of the keys of the dictionary, because the iterator over dictionaries returns its keys (apple, pear, plum) rather than its values (red, green, purple).

fruit_dict = dict(apple='red', pear='green', plum='purple')
print(fruit_dict['apple'])    # 'red'
for fruit in fruit_dict:
    print(fruit)              # 'apple' - 'pear' - 'plum'

Luckily, there are two very powerful functions in Python that can get you the first value of any list, dictionary, generator out there, or indeed the first value of any iterable in the Python universe. Those functions are next() and iter().

next(iter(fruits))              # apple
next(iter(fruit_dict))          # apple
next(iter(fruit_generator))     # apple
next(iter(fruit_defaultdict))   # apple

Amazing! The way it works is straightforward, iter() turns your iterable into an iterator object (“as if you have started a loop”), next() takes the next item of that iterator object (“as if you are in the first iteration of the loop”).

No need to figure out what the keys of the dictionary are, or how to turn the generator into a list. Just two functions and guaranteed success. As an added bonus, it is also very efficient, because almost no memory overhead is created (compared to e.g. list(fruit_dict.keys())[0] which stores all keys of the dictionary in a list.

Finally, sometimes you’re not confident testing just the first element of your list as issues might arise for particular edge cases. No worries, next and iter still have you covered.

temp = iter(fruit_dict)
next(temp)      # apple
next(temp)      # pear
next(temp)      # plum

I often add the following comment to the first line of a loop to make this type of debugging easier, especially if I have to return to code I wrote ages ago.

for fruit in fruit_dict:    # temp = iter(fruit_dict); fruit = next(temp)
    print(fruit)

This allows me to just run the line in the comment and get any example from the iterable I want.

That’s it. Verifying the output of specific examples of your loops is a great way to avoid nasty surprises down the line. Hopefully by making it easier to do this for any kind of iterator, you’ll be able to do it more often!