Pressing Pause On Your Functions: Generators as Control Flow in Python

A major purpose of this post is to improve my understanding of generators–to solidify my understanding of some of the points that James Powell makes in his talk Generators Will Free Your Mind: https://www.youtube.com/watch?v=RdhoN4VVqq8. No credit should be given to me to me for coming up with these ideas. My goal is to articulate them in a different way, in the hopes that that another angle of presenting them will further peoples’ understanding of them.

A common use of generators is for performance, either in regard to CPU or RAM.

# get them all at once
xs = [x%3 for x in range(1_000_000_000)]
# get them one at a time
xs = (x%3 for x in range(1_000_000_000))

If you were around for Python 2, you may remember xrange for this purpose. In Python 2, range returned a list, while xrange returned a generator. This behavior has been changed in Python 3, where xrange no longer exists because range returns a specialized range object that implements the generator protocol.

Using generators to improve performance is completely valid, of course. It usually does improve performance, with little or no impact on the code’s readability or maintainability. It’s a great way to start using generators, and keep using them for as long as you work with Python. But there’s more to them and exploring that is both useful and intellectually rewarding.

Generators can actually serve as a form of control flow. Consider Newton’s Method. The “standard”/textbook implementation of Newton’s Method often takes a number of iterations or an epsilon as a parameter to determine whether it’s converged. Here’s a reference version from Wikipedia’s article on Newton’s method:

def newtons_method(
    x0,               # The initial guess
    f,                # The function whose root we are trying to find
    f_prime,          # The derivative of the function
    tolerance,        # 7-digit accuracy is desired
    epsilon,          # Do not divide by a number smaller than this
    max_iterations,   # The maximum number of iterations to execute
    ):
    for i in range(max_iterations):
        y = f(x0)
        yprime = f_prime(x0)

        if abs(yprime) < epsilon:       # Stop if the denominator is too small
            break

        x1 = x0 - y / yprime            # Do Newton's computation

        if abs(x1 - x0) <= tolerance:   # Stop when the result is within the desired tolerance
            return x1                   # x1 is a solution within tolerance and maximum number of iterations

        x0 = x1                         # Update x0 to start the process again

    return None                         # Newton's method did not converge

There’s nothing at all wrong with this code — it’s good code. It gets the job done, faithfully implements the algorithm, and is very readable. Let’s imagine being the maintainer of a library using this code, and think of how to make it better.

One inconvenience our users may encounter is that they don’t know good values to choose for the guess, tolerance, epsilon, or max iterations. Imagine a user who’s reaching for our numeric library with pip install and isn't very familiar with numeric methods. They have an idea of "good enough," that's clear in their head, but they'd struggle to translate into the parameters of parameters x0, tolerance, epsilon , and max_iterations.

Using a generator and yielding every iterative guess to the caller allows them to decide what “good enough” is on the fly. Here’s what a rough sketch of a Newton’s Method generator may look like:

def newt(f, f_prime, guess):
    while True:
        guess -= f(guess) / f_prime(guess)
        yield guess

This also passes some control back to the caller, since they can investigate each guess and determine whether to continue approximating. That’s incredibly useful for exploratory purposes since the user can run it in an IPython session and investigate the behavior over multiple iterative steps, and also allows library users to wrap dynamic behavior around each iteration, rather than having to wait for the function version to finish aftera number of iterations, determined at initial call time.

This setup also allows users to instrument or add logging however they like. For example:

def harness():
    gen = newt(f, f_prime, 10)
    for i in range(10):
        print(f"{next(gen) = }")

So in this way, generators pass some control back to the user, giving them more flexibility by removing assumptions on how your code is used. I used this concept of generators as a form of flow control in another hobby/personal interest project recently.

Fire Emblem, a Game Boy Advance game I liked as a kid, was recently rereleased on the Nintendo Switch. Fire Emblem is a series of turn-based tactical RPGs. Compared to other turn-based RPGs like Pokémon, most of the game mechanics in Fire Emblem are pretty clear and transparent; for example, battle calculations are shown before a you battle, and since they’re simple arithmetic (damage is character strength rating + weapon strength rating — opponent’s defense rating), a lot of the games’ appeal is that the simplicity of the mechanics adds depth, by giving players a lot of material to use to strategize. You can check your characters’ parameters vs. the CPU’s, figure out how much risk any move opens you up to before making it, and form your strategies using that information.

Gameplay (Combat) in Fire Emblem: The Sacred Stones, © 2004–2005 Nintendo/Intelligent Systems

(As an aside: Games are a great environment to get inspiration for “real” feeling problems, since the rules of a game usually combine both math and rules which emulate the business logic of an application a professional software engineer may work on. While the rules of a game are often quite simplified compared to business logic a, the analogy works alright, as far as being a good environment to lift substantial problems from, that are interesting without feeling too abstract or “homework-y”.)

In Fire Emblem, you and the computer take turns moving the characters in your respective armies around the map to accomplish your goals. Combat either happens on the map, where one character attacks another and risks a counterattack, or in an arena, where two characters take turns attacking until one withdraws or runs out of hit points. I wanted to reuse as much code as possible, so to that end I wanted to define the combat logic in a way that I could easily use the same functions for both types of combat without much awkwardness. Generators helped me here.

Simplifying (the actual code is in a GitHub repository I maintain at https://github.com/EFulmer/python-emblem), the combat function could look a little like this:

def combat(attacker, defender):
    if attack_hits(attacker, defender):
        damage = attacker.strength + attacker.weapon_damage - defender.dfn
        defender.cur_hp -= damage

A naïve formulation of how the Arena combat (infinite rounds until one character is KO’d) could look like this:

def arena(attacker, defender):
    while attacker.cur_hp > 0 and defender.cur_hp > 0:
        attacker, defender = combat(attacker, defender)
        print(f"{attacker.cur_hp = }, {defender.cur_hp = }")
        continue = raw_input("Continue?")
        if continue.upper().trim() == "N":
            break

While this code is rough, it serves a point. This formulation restricts the user in a few ways:

• The arena function repeatedly passes control back to its caller in a nonintuitive way.
• It’s very dependent on environment; a terminal-based game, a desktop GUI game, and a web/mobile game would require their own formulations of the Arena, and all those formulations would repeat the combat rules.

My idea for how to make combat elegant and reusable was to use a generator here, thereby explicitly pausing the execution of combat and making the caller handle it however they felt best. Here’s how that generator formulation looks, more or less:

def arena(attacker, defender):
    while attacker.cur_hp > 0 and defender.cur_hp > 0:
        attacker, defender = combat(attacker, defender)
        defender, attacker = combat(defender, attacker)
        yield attacker, defender
    return attacker, defender  # combat ended because someone ran out of HP

In this formulation, the arena generator explicitly passes control back to the caller and lets them continue gameplay in whatever way they want. If the player wants to end combat, they can simply not call next on the generator. This sort of formulation would allow the game to be extended in a lot of different ways: it could be used in a terminal-based game using something like rich, or it could be put behind a HTTP API and used for a mobile or browser based online game. Passing the flow of gameplay back to the caller simplifies the code and allows the.

This is still a topic I’m exploring myself, so keep an eye out for future blog posts. No promises on more generator-centric content, but I’m certain I’ll play more with it in the future. Thank you for reading!

(Updated on Saturday, 23 September 2023 to reflect a mistake in the code, pointed out to me by my colleague Scott Brownlie)