Python Cheatsheet 05 — Iterators, Generators, and Comprehensions

it = iter([1, 2, 3])
next(it)        # 1
next(it)        # 2
next(it, "done")  # 3 (default on exhaustion)

def count_to(n):
    for i in range(n):
        yield i

for x in count_to(5):
    print(x)

squares = (x**2 for x in range(10))
total = sum(x**2 for x in range(10))    # no list built

# List
squares = [x**2 for x in range(10)]

# Dict
by_id = {u.id: u for u in users}

# Set
unique = {x % 7 for x in numbers}

# Filtered
evens = [x for x in nums if x % 2 == 0]

# Nested
matrix = [[r * c for c in range(5)] for r in range(5)]

async def fetch_pages(urls):
    for u in urls:
        page = await fetch(u)
        yield page

async for page in fetch_pages(urls):
    process(page)

async def all_pages(urls):
    return [page async for page in fetch_pages(urls)]

def parent():
    yield from [1, 2, 3]
    yield from another_gen()

def echo():
    while True:
        msg = yield
        print("got:", msg)

g = echo()
next(g)         # advance to first yield
g.send("hi")    # prints "got: hi"

from itertools import (
    count, cycle, repeat,                    # infinite
    chain, compress, dropwhile, takewhile,    # filtering
    starmap, accumulate, batched,            # transforming
    combinations, permutations, product,      # combinatorial
    groupby, islice, tee, zip_longest,        # iteration
)

# Examples
list(chain([1, 2], [3, 4]))                  # [1, 2, 3, 4]
list(batched("abcdefg", 3))                  # [("a","b","c"), ("d","e","f"), ("g",)]
list(zip_longest([1, 2], [3, 4, 5], fillvalue=0))  # [(1, 3), (2, 4), (0, 5)]

# Pairwise
from itertools import pairwise
list(pairwise([1, 2, 3, 4]))                 # [(1,2), (2,3), (3,4)]

from itertools import islice

first_10 = list(islice(gen, 10))             # take first 10
window = list(islice(gen, 5, 15))            # skip 5, take 10

from itertools import groupby

# Group consecutive identical
items = [("a", 1), ("a", 2), ("b", 3), ("a", 4)]
for key, group in groupby(items, key=lambda x: x[0]):
    print(key, list(group))
# a [(a,1), (a,2)]
# b [(b,3)]
# a [(a,4)]

list(takewhile(lambda x: x < 5, [1, 2, 3, 5, 4]))    # [1, 2, 3]
list(dropwhile(lambda x: x < 5, [1, 2, 3, 5, 4]))    # [5, 4]

from itertools import product

list(product([1, 2], ["a", "b"]))     # [(1,"a"), (1,"b"), (2,"a"), (2,"b")]
list(product(range(2), repeat=3))     # [(0,0,0), (0,0,1), ..., (1,1,1)]

from itertools import chain

nested = [[1, 2], [3, 4], [5]]
flat = list(chain.from_iterable(nested))    # [1, 2, 3, 4, 5]

from collections import Counter

c = Counter("aabbbcccc")     # {'c': 4, 'b': 3, 'a': 2}
c.most_common(2)             # [('c', 4), ('b', 3)]

from collections import defaultdict

by_letter = defaultdict(list)
for word in words:
    by_letter[word[0]].append(word)

for i, item in enumerate(items, start=1):
    print(i, item)

for a, b in zip(list_a, list_b):
    ...

# strict — same length required
for a, b in zip(list_a, list_b, strict=True):
    ...

reversed([1, 2, 3])                           # iterator
sorted(items, key=lambda x: x.name)
sorted(items, key=lambda x: x.age, reverse=True)

# OK
list(map(str, [1, 2, 3]))

# Clearer
[str(x) for x in [1, 2, 3]]