Your service writes a log with a million lines, and you want the first error. This function collects every error line, and you take the first one:
def error_lines(lines: list[str]) -> list[str]:
found: list[str] = []
for line in lines:
if line.startswith("ERROR"):
found.append(line)
return found
log = ["INFO start", "ERROR disk full", "INFO retry", "ERROR disk full"]
print(error_lines(log)[0]) # ERROR disk full
It works, but it checks all million lines and keeps every error in memory to give you one.
Hand back one value at a time with yield
Replace found.append(line) with yield line, and drop the list:
from collections.abc import Iterator
def error_lines(lines: list[str]) -> Iterator[str]:
for line in lines:
if line.startswith("ERROR"):
yield line
log = ["INFO start", "ERROR disk full", "INFO retry", "ERROR disk full"]
errors = error_lines(log)
print(next(errors)) # ERROR disk full
A function with yield in its body is a generator function. Calling it
does not run the body. It returns a generator, an object that runs the body
when you ask it for a value. Each call to next() runs the body up to the next
yield, hands back that value, and pauses there. Here next() stopped at the
second line, so the rest of the log was never checked.
Annotate a generator function with Iterator[...] from collections.abc, and
name the type of the values it yields.
Watch the body run
These print calls show when each part of the body runs:
from collections.abc import Iterator
def count_up(limit: int) -> Iterator[int]:
print("start")
for n in range(1, limit + 1):
yield n
print("end")
numbers = count_up(3) # prints nothing
print(next(numbers)) # prints start, then 1
print(next(numbers)) # prints 2
print(list(numbers)) # prints end, then [3]
print(next(numbers, 0)) # prints 0
When the body ends, the generator has nothing left. A plain next(numbers) would
then raise StopIteration. A for loop and list() catch that for you and
stop. next(numbers, 0) returns its second argument instead of raising.
An iterator works only once
An iterable is anything you can loop over, such as a list, a dict, a string
or a range. An iterator is the object that hands out the values one at a
time and remembers where it is. Each for loop over a list asks the list for a
new iterator, so every loop starts from the beginning. A generator is its own
iterator. Once it reaches the end, it stays empty:
words = ["a", "b"]
print(list(words), list(words)) # ['a', 'b'] ['a', 'b']
shout = (w.upper() for w in words)
print(list(shout), list(shout)) # ['A', 'B'] []
(w.upper() for w in words) is a generator expression. It looks like a list
comprehension in parentheses, and it gives you a generator.
An iterator you can loop over only once is one-shot. Generators, open
files, and the results of map, filter, zip and most of itertools are all
one-shot. They are also lazy, which means they do their work only when
someone asks for the next value.
Common mistakes
Looping twice. A second loop over a used-up generator runs zero times and
raises no error. If you need two passes, build a list once with
rows = list(gen), or call the generator function again to get a fresh
generator.
Expecting a check to run at the call. A check at the top of a generator
body runs at the first next(), not when you call the function:
from collections.abc import Iterator
def pages(size: int) -> Iterator[int]:
if size < 1:
raise ValueError("size must be at least 1")
yield size
bad = pages(0) # no error yet
# next(bad) raises ValueError here
To report a bad argument at the call, check it in a plain function, then return the generator. The batching exercise later in this track asks for exactly that.
Calling len(). A generator does not know how many values it has left, so
len() on one raises TypeError. Build a list first if you need the length.
Check your understanding
You set rows = read_rows(path), where read_rows is a generator function. You count the rows with sum(1 for _ in rows), then loop over rows to save each one. Why is nothing saved?
Show answer
Hide answer
Counting the rows used up the generator. The for loop then got an iterator
that was already at its end, so its body ran zero times. Build a list once with
rows = list(read_rows(path)) if it fits in memory, or call read_rows(path)
again for the second pass.
Going further
A generator can also receive values through send() and finish with
return value. The type Generator[YieldType, SendType, ReturnType] from
collections.abc describes both, and yield from other passes on every value
of another iterable. A later track on generators and streaming data builds on
this article. The Python tutorial covers
generators and
generator expressions.