How a for loop really gets its values, and how yield lets a function hand them out one at a time.
The protocol behind every for loop.
# iterating over a list
numbers = [1, 2, 3, 4, 5]
for num in numbers:
print(num)
# iterating over a dictionary
person = {"name": "Alice", "age": 25}
for key in person:
print(key, person[key])
# iterating over a string
for char in "hello":
print(char)
# iterating over a file
with open("data.txt") as file:
for line in file:
print(line.strip())
Lists, strings, dictionaries and files all speak the same iterator protocol.
# getting an iterator
numbers = [1, 2, 3]
iterator = iter(numbers)
# calling next
next(iterator) # ans: 1
next(iterator) # ans: 2
next(iterator) # ans: 3
next(iterator) # ans: StopIteration
The iterator remembers where it is. The list itself never changes.
for num in numbers:
print(num)
iterator = iter(numbers)
while True:
try:
num = next(iterator)
print(num)
except StopIteration:
break
for calls iter() once, then next() until StopIteration says there is nothing left.
# iterator class
class Counter:
def __init__(self, start, end):
self.current = start
self.end = end
def __iter__(self):
return self
def __next__(self):
if self.current >= self.end:
raise StopIteration
self.current += 1
return self.current - 1
counter = Counter(0, 5)
for num in counter:
print(num) # ans: 0, 1, 2, 3, 4
# manual iteration
counter = Counter(0, 3)
next(counter) # ans: 0
next(counter) # ans: 1
next(counter) # ans: 2
next(counter) # ans: StopIteration
has __iter__(): you can ask it for an iterator
has __iter__() and __next__(): it hands out values
# a list is iterable but not an iterator
numbers = [1, 2, 3]
iter(numbers) # creates an iterator
next(numbers) # ans: Error! (TypeError)
# an iterator is both iterable and iterator
iterator = iter(numbers)
iter(iterator) # ans: returns itself
next(iterator) # ans: 1
# range is iterable
r = range(5)
iter(r) # creates an iterator
next(r) # ans: Error! (TypeError)
class MyRange:
def __init__(self, start, end):
self.start = start
self.end = end
def __iter__(self):
return MyRangeIterator(
self.start, self.end)
class MyRangeIterator:
def __init__(self, start, end):
self.current = start
self.end = end
def __iter__(self):
return self
def __next__(self):
if self.current >= self.end:
raise StopIteration
self.current += 1
return self.current - 1
Every iter() call builds a fresh iterator, so one MyRange can be looped over again and again.
Functions that pause at yield and pick up where they left off.
# generator function with yield
def count_up_to(n):
count = 1
while count <= n:
yield count
count += 1
gen = count_up_to(3)
next(gen)
next(gen)
next(gen)
next(gen)
# creates a generator object
gen = count_up_to(5)
next(gen) # ans: 1
next(gen) # ans: 2
# iterate over a generator
for num in count_up_to(5):
print(num) # ans: 1, 2, 3, 4, 5
# simple generator
def simple_gen():
yield 1
yield 2
yield 3
for value in simple_gen():
print(value) # ans: 1, 2, 3
# generator with multiple yields
def fibonacci(n):
a, b = 0, 1
for _ in range(n):
yield a
a, b = b, a + b
list(fibonacci(10))
# ans: [0, 1, 1, 2, 3, 5, 8, 13, 21, 34]
def infinite_counter():
count = 0
while True:
yield count
count += 1
# use with break, or take the first n
counter = infinite_counter()
for i, num in enumerate(counter):
print(num)
if i >= 4:
break
# ans: 0, 1, 2, 3, 4
while True is safe here: the generator only runs when someone asks for the next value.
def even_numbers(n):
for i in range(n):
if i % 2 == 0:
yield i
list(even_numbers(10)) # ans: [0, 2, 4, 6, 8]
# like a list comprehension, in round brackets
gen = (x**2 for x in range(10))
next(gen) # ans: 0
next(gen) # ans: 1
# vs list comprehension
list_comp = [x**2 for x in range(10)] # ans: creates a list
gen_exp = (x**2 for x in range(10)) # ans: creates a generator
import sys
sys.getsizeof([x for x in range(1000)]) # ans: ~9000 bytes
sys.getsizeof((x for x in range(1000))) # ans: ~200 bytes
The list holds all 1000 values. The generator holds only its recipe and where it is up to.
# use in functions: no extra brackets needed
sum(x**2 for x in range(10)) # ans: 285
max(x for x in range(10) if x % 2 == 0) # ans: 8
# chaining generators
numbers = (x for x in range(10))
evens = (x for x in numbers if x % 2 == 0)
squares = (x**2 for x in evens)
list(squares) # ans: [0, 4, 16, 36, 64]
# yield from (Python 3.3+)
def generator1():
yield 1
yield 2
def generator2():
yield 3
yield 4
def combined():
yield from generator1()
yield from generator2()
list(combined()) # ans: [1, 2, 3, 4]
def combined():
for value in generator1():
yield value
for value in generator2():
yield value
# flatten nested structure
def flatten(nested):
for item in nested:
if isinstance(item, list):
yield from flatten(item)
else:
yield item
nested = [1, [2, [3, 4], 5], 6]
list(flatten(nested)) # ans: [1, 2, 3, 4, 5, 6]
send(), throw() and close(): talking back to a paused generator.
def echo():
while True:
value = yield
print(f"Received: {value}")
gen = echo()
next(gen) # prime the generator
gen.send("Hello") # ans: Received: Hello
gen.send("World") # ans: Received: World
next() runs it to the first yield; send(x) resumes it, and yield evaluates to x.
def running_average():
total = 0
count = 0
average = None
while True:
value = yield average
total += value
count += 1
average = total / count
avg = running_average()
next(avg) # ans: None
avg.send(10) # ans: 10.0
avg.send(20) # ans: 15.0
avg.send(30) # ans: 20.0
def my_gen():
try:
yield 1
yield 2
yield 3
except ValueError:
yield "Error caught"
yield 4
gen = my_gen()
next(gen) # ans: 1
gen.throw(ValueError) # ans: "Error caught"
next(gen) # ans: 4
def count_forever():
count = 0
try:
while True:
yield count
count += 1
finally:
print("Cleanup")
gen = count_forever()
next(gen) # ans: 0
next(gen) # ans: 1
gen.close() # ans: prints Cleanup
next(gen) # ans: StopIteration
Ready-made iterators: infinite, combinatoric and utilities.
from itertools import count, cycle, repeat
# count: infinite counting
for i in count(10, 2): # start=10, step=2
print(i)
if i > 20:
break
# ans: 10, 12, 14, 16, 18, 20, 22
# repeat: repeat a value
for x in repeat(10, 3):
print(x)
# ans: 10, 10, 10
The if runs after print, so 22 is printed before the loop breaks.
# cycle: infinite cycling
colors = ['red', 'green', 'blue']
color_cycle = cycle(colors)
for i, color in enumerate(color_cycle):
print(color)
if i >= 5:
break
# ans: red, green, blue, red, green, blue
from itertools import product, permutations
# product: cartesian product
list(product([1, 2], ['a', 'b']))
# ans: [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b')]
# permutations: all orderings
list(permutations([1, 2, 3], 2))
# ans: [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]
from itertools import combinations, combinations_with_replacement
# combinations: unique selections
list(combinations([1, 2, 3, 4], 2))
# ans: [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]
# combinations_with_replacement
list(combinations_with_replacement([1, 2, 3], 2))
# ans: [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]
from itertools import chain, islice, zip_longest
# chain: combine iterables
list(chain([1, 2], [3, 4], [5, 6]))
# ans: [1, 2, 3, 4, 5, 6]
# islice: slice an iterator (start, stop, step)
list(islice(range(10), 2, 8, 2))
# ans: [2, 4, 6]
# zip_longest: zip with a fill value
list(zip_longest([1, 2], ['a', 'b', 'c'], fillvalue=0))
# ans: [(1, 'a'), (2, 'b'), (0, 'c')]
# groupby: group consecutive items
from itertools import groupby
data = [1, 1, 2, 2, 2, 3, 1, 1]
for key, group in groupby(data):
print(key, list(group))
# ans: 1 [1, 1]
# ans: 2 [2, 2, 2]
# ans: 3 [3]
# ans: 1 [1, 1]
Only consecutive equal items are grouped, so 1 appears twice.
Big files, batches, pipelines and coroutines.
def read_large_file(file_path):
with open(file_path) as file:
for line in file:
yield line.strip()
for line in read_large_file("large_file.txt"):
process(line)
Only one line is in memory at a time, however large the file.
def batch_iterator(iterable, batch_size):
batch = []
for item in iterable:
batch.append(item)
if len(batch) >= batch_size:
yield batch
batch = []
if batch:
yield batch
for batch in batch_iterator(range(10), 3):
print(batch)
# ans: [0, 1, 2]
# ans: [3, 4, 5]
# ans: [6, 7, 8]
# ans: [9]
def read_data():
for i in range(100):
yield i
def filter_evens(numbers):
for num in numbers:
if num % 2 == 0:
yield num
def square(numbers):
for num in numbers:
yield num ** 2
pipeline = square(filter_evens(read_data()))
result = list(pipeline)
Each value travels the whole pipeline before the next one is read.
# memory efficient
def read_numbers_list(n):
return [i for i in range(n)]
def read_numbers_gen(n):
for i in range(n):
yield i
# memory usage
import sys
n = 1000000
list_size = sys.getsizeof(read_numbers_list(n))
gen_size = sys.getsizeof(read_numbers_gen(n))
# ans: list_size: ~8MB
# ans: gen_size: ~200 bytes
from itertools import islice
def expensive_operation(x):
# simulate expensive computation
return x ** 2
# generator: operations done on demand
gen = (expensive_operation(x) for x in range(1000000))
first_five = list(islice(gen, 5))
# list: all operations done immediately
lst = [expensive_operation(x) for x in range(1000000)]
The generator calls expensive_operation 5 times; the list calls it a million times.
# simple coroutine
def grep(pattern):
print(f"Searching for {pattern}")
while True:
line = yield
if pattern in line:
print(line)
search = grep("python")
next(search) # prime the coroutine
search.send("I love python") # ans: I love python
search.send("Hello world") # ans: (no output)
search.send("python is great") # ans: python is great
def producer():
for i in range(10):
yield i
def filter_even(target):
while True:
num = yield
if num % 2 == 0:
target.send(num)
def printer():
while True:
value = yield
print(value)
# set up the pipeline
p = printer()
next(p)
f = filter_even(p)
next(f)
for num in producer():
f.send(num)
# ans: 0, 2, 4, 6, 8
A generator pulls values with next(); a coroutine has values pushed into it with send().
has __iter__(); iter(x) gives an iterator
has __next__(); next() until StopIteration
hands out a value and pauses the function, locals and all
a generator expression: lazy, tiny in memory
hands out every value of another iterable
push a value, an exception or a stop into a paused generator