Comprehensions & Generators

Python · 8 entries

List Comprehensions (Detailed)

Syntax

[expr for x in iterable if cond]

Example

# Transform + filter
temps_f = [32, 68, 95, 104, 50]
warm_c = [(f - 32) * 5 / 9 for f in temps_f if f > 60]
print([round(c, 1) for c in warm_c])

# Nested comprehension (flatten)
matrix = [[1, 2], [3, 4], [5, 6]]
flat = [n for row in matrix for n in row]
print(flat)

Output

[20.0, 35.0, 40.0]
[1, 2, 3, 4, 5, 6]

Note In nested comprehensions, the outer loop comes first: [x for row in matrix for x in row]. This matches the order of equivalent nested for loops.

list comprehensionflatten listnested comprehensionfilter and transform

Dict Comprehensions (Detailed)

Syntax

{key_expr: val_expr for item in iterable if cond}

Example

# Invert a dictionary
http_codes = {200: "OK", 404: "Not Found", 500: "Server Error"}
code_lookup = {msg: code for code, msg in http_codes.items()}
print(code_lookup["Not Found"])

# From two parallel lists
keys = ["host", "port", "debug"]
vals = ["localhost", 8080, True]
config = {k: v for k, v in zip(keys, vals)}
print(config)

Output

404
{'host': 'localhost', 'port': 8080, 'debug': True}

Note When inverting a dict, duplicate values in the original will cause key collisions; the last one wins.

dict comprehensioninvert dictionaryzip to dictbuild dict

Set Comprehensions

Syntax

{expr for item in iterable if cond}

Example

sentence = "the quick brown fox jumps over the lazy dog"
word_lengths = {len(word) for word in sentence.split()}
print(sorted(word_lengths))

Output

[3, 4, 5]

Note Set comprehensions look like dict comprehensions but without the colon. They automatically deduplicate results.

set comprehensionunique values comprehensiondeduplicate transform

Generator Expressions

Syntax

(expr for item in iterable if cond)

Example

# Sum of squares without building a list
total = sum(n ** 2 for n in range(1, 101))
print(total)

# Lazy evaluation
import sys
list_size = sys.getsizeof([x for x in range(10000)])
gen_size = sys.getsizeof(x for x in range(10000))
print(f"List: {list_size} bytes, Generator: {gen_size} bytes")

Output

338350
List: ~85000 bytes, Generator: ~200 bytes

Note Generator expressions use parentheses instead of brackets and are lazy - they produce values on demand. When passed as the sole argument to a function, extra parens can be omitted.

generator expressionlazy listmemory efficientgenexp

yield from (Delegation)

Syntax

yield from iterable

Example

def flatten(nested):
    for item in nested:
        if isinstance(item, list):
            yield from flatten(item)
        else:
            yield item

data = [1, [2, [3, 4]], [5, 6]]
print(list(flatten(data)))

Output

[1, 2, 3, 4, 5, 6]

Note yield from delegates to a sub-generator and passes values through. It handles send() and throw() transparently, unlike a manual for-loop with yield.

yield fromdelegate generatorflatten recursivesub-generator

itertools Highlights

Syntax

import itertools

Example

from itertools import chain, islice, groupby, product

# Chain multiple iterables
all_items = list(chain([1, 2], [3, 4], [5]))
print(all_items)

# Take first N from any iterable
first_three = list(islice(range(1000000), 3))
print(first_three)

# Cartesian product
sizes = ["S", "M"]
colors = ["red", "blue"]
combos = list(product(sizes, colors))
print(combos)

Output

[1, 2, 3, 4, 5]
[1, 2, 3]
[('S', 'red'), ('S', 'blue'), ('M', 'red'), ('M', 'blue')]

Note itertools functions return lazy iterators. They compose well and avoid creating intermediate lists. See also: accumulate, starmap, tee, zip_longest.

itertoolschainisliceproductgroupby combinations

itertools.groupby

Syntax

from itertools import groupby
for key, group in groupby(iterable, key=func):

Example

from itertools import groupby

transactions = [
    {"type": "credit", "amount": 100},
    {"type": "credit", "amount": 200},
    {"type": "debit", "amount": 50},
    {"type": "debit", "amount": 75},
]

# Must be sorted by the grouping key first!
for ttype, group in groupby(transactions, key=lambda t: t["type"]):
    amounts = [t["amount"] for t in group]
    print(f"{ttype}: {amounts}")

Output

credit: [100, 200]
debit: [50, 75]

Note groupby only groups consecutive items with the same key. Sort the data by the key first, or you get fragmented groups.

groupbygroup itemscategorizecluster by key

Walrus Operator in Comprehensions

Syntax

[y := f(x) ... if y > threshold]

Example

import math

values = [1, 10, 100, 1000, 10000]
# Compute once, use in both filter and output
result = [
    (v, log_v)
    for v in values
    if (log_v := math.log10(v)) >= 2
]
print(result)

Output

[(100, 2.0), (1000, 3.0), (10000, 4.0)]

Note The walrus operator avoids computing the same expression twice in a comprehension. The variable leaks into the enclosing scope, which can be surprising.

walrus comprehensionassignment in comprehensionavoid double computation