Generators: Solutions

If you read the total as the result of the last send(), writing total: Total = t.send(Amount(12)) with no try, the script prints the three prompts and then stops with StopIteration: 42. ty rejects that line before it runs: send() returns the YieldType, so it reports an invalid-assignment of a Prompt to a Total. A generator’s return value arrives only on its StopIteration, so the solution catches that exception and reads its value.

# exercise_1.py
from collections.abc import Generator
from typing import NewType

Prompt = NewType("Prompt", str)
Amount = NewType("Amount", int)
Total = NewType("Total", int)

def tally() -> Generator[Prompt, Amount, Total]:
    total = 0
    for n in (1, 2, 3):
        amount = yield Prompt(f"amount {n} of 3")
        total += amount
    return Total(total)

t = tally()
print(next(t))
#: amount 1 of 3
print(t.send(Amount(10)))
#: amount 2 of 3
print(t.send(Amount(20)))
#: amount 3 of 3
try:
    t.send(Amount(12))
except StopIteration as stop:
    total: Total = stop.value
print(total)
#: 42

Keep the three channels apart. The three-parameter annotation names all three channels: Generator[Prompt, Amount, Total] says this generator yields a Prompt, receives an Amount, and finally returns a Total. Three NewType definitions over str, int, and int keep the two integer channels apart, so transposing the SendType and the ReturnType is a type checker error rather than a bug that shows up in arithmetic.

Drive the conversation by hand. Driving tally() by hand takes four calls: one next() and three sends. next(t) runs the body up to the first yield and produces the first prompt. Each send() resumes at that suspended yield, whose value becomes amount, then runs to the next one. The third send() finds no fourth yield, so the loop ends and tally() returns. The return value arrives as StopIteration’s value rather than as the result of send().

total is the generator’s own local, and it survives across three suspensions with no storage anywhere else. The frame is the state.

# exercise_2.py
from collections.abc import Generator, Iterator
from typing import NewType

Question = NewType("Question", str)
Answer = NewType("Answer", str)
Result = NewType("Result", str)

def interview() -> Generator[Question, Answer, Result]:
    name = yield Question("name")
    town = yield Question("town")
    friend = yield Question("friend")
    return Result(f"{name} of {town}, friend {friend}")

def drive_from_dict(
        conversation: Generator[Question, Answer, Result],
        answers: dict[Question, Answer]) -> Result:
    request = next(conversation)
    while True:
        answer = answers[request]
        print(f"{request = }, {answer = }")
        try:
            request = conversation.send(answer)
        except StopIteration as stop:
            return stop.value

def drive_naive(
        conversation: Generator[Question, Answer, Result],
        answers: Iterator[Answer]) -> Result:
    request = next(conversation)
    while True:
        try:  # Both next() calls share one except clause
            reply = next(answers)
            print(f"{request = }, {reply = }")
            request = conversation.send(reply)
        except StopIteration as stop:
            return stop.value

def drive_in_order(
        conversation: Generator[Question, Answer, Result],
        answers: Iterator[Answer]) -> Result:
    request = next(conversation)
    while True:
        # Fetched outside the try on purpose
        reply = next(answers)
        print(f"{request = }, {reply = }")
        try:
            request = conversation.send(reply)
        except StopIteration as stop:
            return stop.value

by_name = {
    Question("name"): Answer("Alice"),
    Question("town"): Answer("Wonderland"),
    Question("friend"): Answer("Rabbit"),
}
print(drive_from_dict(interview(), by_name))
#: request = 'name', answer = 'Alice'
#: request = 'town', answer = 'Wonderland'
#: request = 'friend', answer = 'Rabbit'
#: Alice of Wonderland, friend Rabbit
in_order = iter([Answer("Alice"), Answer("Wonderland"),
                 Answer("Rabbit")])
print(drive_in_order(interview(), in_order))
#: request = 'name', reply = 'Alice'
#: request = 'town', reply = 'Wonderland'
#: request = 'friend', reply = 'Rabbit'
#: Alice of Wonderland, friend Rabbit
# One answer, three questions:
try:
    drive_in_order(interview(), iter([Answer("Alice")]))
except StopIteration:
    print("answer source ran out")
#: request = 'name', reply = 'Alice'
#: answer source ran out
print(repr(drive_naive(interview(),
                       iter([Answer("Alice")]))))
#: request = 'name', reply = 'Alice'
#: None

Reuse the generator unchanged. Nothing in interview() changes, and nothing could have. It yields a Question and receives an Answer without asking about the answer’s source. That is the separation the chapter teaches: the generator describes the conversation, and the driver interprets it. Swapping one interpreter for another leaves the description untouched.

Find the answer for each request. The two drivers differ in the property on which they rely. The dictionary driver looks each answer up by the request, so it answers correctly in whatever order the questions arrive, and it answers a repeated question the same way twice. The iterator driver goes by position, so it depends on the generator asking the questions for which the driver has replies, in that order. Both satisfy the same type. The type says what travels, not what the driver knows.

Keep the two endings apart. One detail in drive_in_order() earns its comment. next(answers) sits outside the try because the except StopIteration meant for the conversation otherwise catches a StopIteration raised by an exhausted answer list. Two different iterators raising one exception type is a real hazard when a driver holds both.

Test a short answer source. The last two runs show that hazard. Given one answer and three questions, drive_in_order() lets the StopIteration escape, so the caller learns the answer source ran dry. drive_naive() differs only in having next(answers) inside the try. It catches that same exception, reads it as “the conversation finished,” and returns stop.value, which is None.

None is the wrong answer twice over. The interview did not finish, so no Result exists, and None is not a Result in any case. Nothing catches the mistake: StopIteration.value has type Any, so return stop.value satisfies a declared Result and the checker reports nothing. The failure is silent at the type checker and silent at runtime. It surfaces later as a None where the caller expects a string, far from the driver that produced it.

Keeping the two meanings apart is a one-line discipline: put inside the try only the call whose StopIteration you mean to interpret.

# exercise_3.py
from collections.abc import Generator

def collect(name: str) -> Generator[str, int]:
    first = yield f"{name} needs a value"
    second = yield f"{name} needs another"
    print(f"{name} got {first} and {second}")

def both() -> Generator[str, int]:
    yield from collect("alpha")
    yield from collect("beta")
    yield from collect("gamma")

g = both()
print(next(g))
#: alpha needs a value
for value in [1, 2, 3, 4, 5]:
    print(g.send(value))
#: alpha needs another
#: alpha got 1 and 2
#: beta needs a value
#: beta needs another
#: beta got 3 and 4
#: gamma needs a value
#: gamma needs another
try:
    g.send(6)
except StopIteration:
    print("both() is exhausted")
#: gamma got 5 and 6
#: both() is exhausted

Count the values the collectors need. The five sends do not divide evenly among three collectors. Each collect() consumes two values, so the three collectors need six, and the loop supplies five. gamma stays suspended at its second yield until the send(6) after the loop completes it.

Read the output in pairs. The pairs of lines are the tell. A send that supplies a collector’s first value produces one line, the same collector’s second prompt. A send that supplies a collector’s second value produces two: the completed collector’s print(), then the first prompt of the next collector. The two-line sends are send(2), send(4), and send(6), so the output alternates between one-line and two-line responses all the way down. The second line of send(6)’s pair is the exception: no collector remains to prompt, so gamma finishes, both() raises StopIteration, and the except prints both() is exhausted.

both() takes no part in that alternation. It contains three yield from statements and no code that forwards a value. yield from relays in both directions on its own: prompts up to the driver, numbers down to whichever yield is currently suspended, two frames below.

If you change report()’s annotation but leave out return size, the script still runs, and the printed list ends with 'total: None': report() returns None, and the second yield from delivers that None into counted. ty reports an invalid-return-type on the annotation, because with no return statement report() always implicitly returns None, and its annotation declares an int return value. The solution returns size from report(), so the count reaches summarize() through a second return.

# exercise_5.py
from collections.abc import Generator

def emit(items: list[str]) -> Generator[str, None, int]:
    total = 0
    for item in items:
        yield item
        total += len(item)
    return total

def report(items: list[str]) -> Generator[str, None, int]:
    size: int = yield from emit(items)
    yield f"({size} characters)"
    return size

def summarize(items: list[str]) -> Generator[str]:
    counted: int = yield from report(items)
    yield f"total: {counted}"

print(list(summarize(["red", "green", "blue"])))
#: ['red', 'green', 'blue', '(12 characters)', 'total: 12']

Declare the return channel. report()’s annotation changes from Iterator[str] to Generator[str, None, int], because a generator that returns something needs the long form. Iterator sets the ReturnType to None, so ty rejects report()’s own return size with expected None, found int.

Keep the strings and the count apart. The strings and the count travel by different channels, and the listing shows both at once. Every string that emit() or report() yields travels through the YieldType and comes out in the list. The count travels through the ReturnType: emit() returns it, yield from delivers it into report()’s size, report() returns it again, and the second yield from delivers it into summarize()’s counted. The SendType is None throughout, since nobody sends anything in.

12 therefore reaches summarize() through two returns and is not yielded on its own. '(12 characters)' takes the other route: report() yields it, the yield from in summarize() relays it, and the driver receives it. No generator binds it to a name. The same number can travel either way, and the choice decides who can see it: a yielded value goes to the driver, a returned value goes to the delegating generator.

If you write stock = STOCK instead of copying with dict(STOCK), each sale changes the module-level table from which every machine starts. After three "12" sales, by any mix of machines, a fresh machine() answers the next "12" order with UNAVAILABLE. The Final annotation forbids rebinding the name, not changing the dictionary, so the type checker passes the shared version. The solution copies STOCK into a local, so each generator’s frame holds its own stock.

# exercise_7.py
from collections.abc import Generator
from typing import Final, Literal
from record import record

@record
class Coin:
    cents: int

@record
class Digit:
    value: str

type Event = Coin | Digit
type Report = Literal[
    "QUIESCENT", "COLLECTING", "SELECTING",
    "UNAVAILABLE", "WANT_MORE", "DISPENSED"]

PRICES: Final[dict[str, int]] = {"11": 25, "12": 75}
STOCK: Final[dict[str, int]] = {"11": 0, "12": 3}

def machine() -> Generator[Report, Event]:
    stock = dict(STOCK)
    amount = 0
    event: Event = yield "QUIESCENT"
    while True:
        while isinstance(event, Coin):
            amount += event.cents
            event = yield "COLLECTING"
        row = event.value  # Not a Coin, so a first Digit
        second = yield "SELECTING"
        # A coin instead of a digit
        if isinstance(second, Coin):
            amount += second.cents
            event = yield "COLLECTING"
            continue
        code = row + second.value
        if stock.get(code, 0) == 0:
            event = yield "UNAVAILABLE"
        elif amount < PRICES.get(code, 0):
            event = yield "WANT_MORE"
        else:
            amount -= PRICES[code]
            stock[code] -= 1
            event = yield "DISPENSED"

m = machine()
print(next(m))
#: QUIESCENT
for event in [Coin(25), Digit("1"), Digit("1"), Digit("1"),
              Digit("2"), Coin(50), Digit("1"), Digit("2")]:
    print(f"{event} -> {m.send(event)}")
#: Coin(cents=25) -> COLLECTING
#: Digit(value='1') -> SELECTING
#: Digit(value='1') -> UNAVAILABLE
#: Digit(value='1') -> SELECTING
#: Digit(value='2') -> WANT_MORE
#: Coin(cents=50) -> COLLECTING
#: Digit(value='1') -> SELECTING
#: Digit(value='2') -> DISPENSED

Let position carry the state. The machine holds no state attribute and consults no table. Where the generator pauses is the state: paused in the coin loop means COLLECTING, paused after yield "SELECTING" means a first digit has arrived and the machine waits for a second. amount, row, and stock are locals that survive because the frame does. This version has no counterpart for two parts of the table-driven version: the state attribute and the transition lookup.

Report the state reached. The yield in machine() runs the opposite direction from the one in interview(), and neither signature says so. Both yield strings, but interview() yields a request the driver must satisfy, while machine() yields a report the driver may ignore. The driver’s event and the machine’s report travel independently: send(Coin(25)) answers no question, it delivers an event. A generator’s type describes the traffic, not who is in charge. Both arrangements fit the same Generator annotation.

For another state, take the table. The generator’s compactness comes from the states forming a line, so control flow can express the sequence. The two states here that break the line cost something: an if chain reaches UNAVAILABLE and WANT_MORE, and each one returns by looping back to the top, a goto written as a while True. Now add a state reachable from three others, the way the table handles Quit from every state but QUIESCENT. No position in the body corresponds to it. The new state becomes a flag, or a check repeated at several yields, and either one breaks the correspondence between position and state, the one thing that makes this version readable.

The table pays a fixed cost instead. Adding a state means one new Enum member and a few new rows. Those rows sit next to the existing ones, where you can read the whole machine at once. The generator is the better choice for a conversation with a beginning and an end, like interview(). The table is the better choice for a machine that runs forever and can go anywhere from anywhere.