6 Tips To Write Better For Loops in Python

6 Tips To Write Better For Loops in Python

In this tutorial we'll have a look at 6 tips to write better For-Loops in Python. These tips include easy refactoring methods that you can apply right away in your code.

1) Don't use loops at all¶

If possible, try to avoid loops and use built-in methods instead.

#more

A typical example is to use the built-in sum() method instead of manually looping:

numbers = [10, 20, 33, 40, 50, 60, 70, 80]
result = 0
for num in numbers:
    result += num

print(result) # 363

# --> Refactor
result = sum(numbers)
print(result) # 363

2) Use enumerate¶

When you need to iterate over a list, and you need both the index and the current item, use the built-in enumerate() function instead:

data = ["a", "b", "c"]
for i in range(len(data)):
    print(i, data[i])


# --> Refactor
for idx, val in enumerate(data):
    print(idx, val)

# 0 a
# 1 b
# 2 c

It takes an optional start parameter that starts index counting at this value.

for idx, val in enumerate(data, start=1):
    print(idx, val)

# 1 a
# 2 b
# 3 c

3) Use zip¶

When iterating over multiple lists, use zip. It returns an iterator of tuples, and stops when the shortest input iterable is exhausted:

a = [1, 2, 3]
b = ["a", "b", "c"]

for i in range(len(a)):
    print(a[i], b[i])


# --> Refactor
for val1, val2 in zip(a, b):
    print(val1, val2)

# 1 a
# 2 b
# 3 c

Since Python 3.10, you can set the parameter strict=True to raise errors if one of the iterables is shorter than the other. This might be helpful to catch bugs in your program.

a = [1, 2, 3]
b = ["a", "b", "c", "d"]

for val1, val2 in zip(a, b, strict=True):
    print(val1, val2)

# ValueError: zip() argument 2 is longer than argument 1

4) Think lazy!¶

Try to think lazy and refactor your code with generators.

I found this great tip in this Python talk: A deeper look at iteration in Python

Here, we use a generator expression that can then be assigned to a variable, helping to make the code more readable. When the generator object is then used with the sum() function, we get the same result. Only at this point calculations are performed.

events = [("learn", 5), ("learn", 10), ("relaxed", 20)]

minutes_studied = 0
for event in events:
    if event[0] == "learn":
        minutes_studied += event[1]

print(minutes_studied) # 15

# --> Refactor
study_times = (event[1] for event in events if event[0] == "learn")
minutes_studied = sum(study_times)
print(minutes_studied) # 15

5) Use itertools more!¶

The itertools module provides functions creating iterators for efficient looping. Let's look at three of those functions:

• islice()
• pairwise()
• takewhile()

itertools.islice()¶

It makes an iterator that returns selected elements from the iterable.

lines = ["line1", "line2", "line3", "line4", "line5",
         "line6", "line7", "line8", "line9", "line10"]

for i, line in enumerate(lines):
    if i >= 5:
        break
    print(line)

# --> Refactor
from itertools import islice

first_five_lines = islice(lines, 5)
for line in first_five_lines:
    print(line)

# line1
# line2
# line3
# line4
# line5

pairwise()¶

Returns successive overlapping pairs taken from the input iterable.

 data = 'ABCDE'
for i in range(len(data)-1):
    print(data[i], data[i+1])

# --> Refactor
from itertools import pairwise
for pair in pairwise('ABCDE'):
    print(pair[0], pair[1])

# A B
# B C
# C D
# D E

takewhile()¶

It makes an iterator that returns elements from the iterable as long as the predicate is true.

for item in [1, 2, 4, -1, 4, 1]:
    if item >= 0:
        print(item)
    else:
        break

# --> Refactor    
from itertools import takewhile
items = takewhile(lambda x: x >= 0, [1, 2, 4, -1, 4, 1])
for item in items:
    print(item)

# 1
# 2
# 4

6) Use NumPy¶

If speed is very important, you can use NumPy. It provides many corresponding functions to built-in functions like np.sum(), np.arange() etc.

NumPy functions are usually much faster. However, be aware that with functions like np.arange() the whole array is allocated in memory and thus takes more space as a result, while Python's range() function doesn't need to allocate a whole array.

import numpy as np

sum(range(10))

# --> Refactor  
np.sum(np.arange(10))

Further Resources:¶

• Talk: a deeper look at iteration in Python