Python comprehensions, generators and decorators

List comprehensions: build a list in one line

squares = []
for n in range(1, 6):
    squares.append(n * n)

squares = [n * n for n in range(1, 6)]       # the same, in one line

The pattern is [expression for item in iterable if condition]:

marks = [45, 82, 91, 30, 67]
passed = [m for m in marks if m >= 33]            # [45, 82, 91, 67]
names = [s.title() for s in ["asha", "ravi"]]     # ['Asha', 'Ravi']
labels = ["Pass" if m >= 33 else "Fail" for m in marks]

Note where the if goes: a filter goes at the end; a choice between two values goes at the start.

Dict and set comprehensions

names = ["Asha", "Ravi", "Meena"]
lengths = {n: len(n) for n in names}         # {'Asha': 4, 'Ravi': 4, 'Meena': 5}
prices = {"pen": 10, "bag": 450}
with_gst = {k: v * 1.18 for k, v in prices.items()}
initials = {n[0] for n in names}             # a set: {'A', 'R', 'M'}

table = [[r * c for c in range(1, 4)] for r in range(1, 4)]   # nested: a 3x3 grid

Keep comprehensions short. If one does not fit comfortably on a line, a normal loop is clearer.

map, filter, zip, enumerate, any and all

names = ["Asha", "Ravi"]
marks = [91, 78]
for name, m in zip(names, marks):
    print(f"{name}: {m}")

A common use of map: read several numbers on one line.

a, b, c = map(int, input("Three numbers: ").split())

sorted(), min() and max() with a key

The key function tells Python *what to compare*:

students = [("Asha", 91), ("Ravi", 78), ("Meena", 85)]
by_marks = sorted(students, key=lambda s: s[1], reverse=True)
topper = max(students, key=lambda s: s[1])
shortest = min(["Ranchi", "Gaya", "Patna"], key=len)    # 'Gaya'

Iterators: how a for loop really works

A for loop calls iter() on what you give it, then next() again and again until there are no items left:

it = iter(["a", "b"])
print(next(it))     # a
print(next(it))     # b
next(it)            # StopIteration: nothing left

Generators: functions that yield values one at a time

A generator function uses yield instead of return. Each yield hands back one value and pauses the function until the next value is asked for.

def countdown(n):
    while n > 0:
        yield n
        n -= 1

for x in countdown(3):
    print(x)          # 3, 2, 1

A generator expression looks like a list comprehension with round brackets:

total = sum(n * n for n in range(1_000_000))

Generators save memory

A list holds all its values at once; a generator makes them one at a time. To read a huge log file, a generator handles one line at a time instead of loading millions of lines into memory:

def errors(path):
    with open(path, encoding="utf-8") as f:
        for line in f:
            if "ERROR" in line:
                yield line.strip()

for e in errors("server.log"):
    print(e)

Decorators: wrap a function to add a feature

A decorator takes a function and returns a new function that does something extra before or after calling the original. You apply it with @.

import time

def timer(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        print(f"{func.__name__} took {time.time() - start:.3f} s")
        return result
    return wrapper

@timer                       # same as: slow_sum = timer(slow_sum)
def slow_sum(n):
    return sum(range(n))

slow_sum(10_000_000)

Real apps use decorators for login checks (@login_required in Flask and Django), caching (@functools.lru_cache) and logging.

Common mistakes

• Putting a filter if before the for in a comprehension.
• Expecting map() or filter() to give a list: wrap them in list() to see the values.
• Looping over a generator twice: it is used up after the first loop.
• Forgetting to return result in a decorator's wrapper, so the decorated function returns None.

Key points

• [expr for x in items if cond] builds a list; {k: v for ...} and {x for ...} build dicts and sets.
• zip, enumerate, any and all replace many hand-written loops.
• key= tells sorted, min and max what to compare.
• Generators yield values one at a time; decorators wrap functions with @.

Practice questions

def evens(n):
    for i in range(n):
        yield 2 * i