Waltzing with Python’s Walrus Operator

Using the walrus operator := to make life simpler.

Python 3.8 introduced a new assignment operator with PEP 572 called assignment expressions, a.k.a the walrus operator. The walrus operator uses the new walrus-like syntax :=, to assign variables within an expression.

It’s been out for a few years at this point (at the time of writing Python 3.12 is around the corner) and I’ve found some joy in how it’s helped elegantly shorten some parts of my code. Here are the ways I’ve made use of the walrus operator.

Error Handling

The walrus operator can help reduce repetition and make error handling a bit more streamlined. In the following snippet, func() will return None to represent an error occurred.

x = func()
if not x:
    print("Error message")
    return

Using the walrus operator, the call to func() can be inlined.

if not x := func():
    print("Error message")
    return

Shaving off a single line may seem trivial but those saved lines can add up. For example, when parsing user inputs and performing validation. In the following snippet, we want to validate user inputs x, y, and z. If there’s a validation problem, validate will return a string with a message explaining what is wrong with the input and a message of how to fix. These messages get appended to a list so all validation messages can be printed out together.

validation_errors = []
if msg := validate(x):
    validation_errors += msg

if msg := validate(y):
    validation_errors += msg

if msg := validate(z):
    validation_errors += msg

if validation_errors: # a non-empty list resolves to True
    for error in validation_errors:
        print(error)
    return

Comprehensions

Let’s say we wanted to create a list of results from expensive function call but only results that aren’t None. With a list comprehension, the expensive function would need to be called twice. Not ideal.

y = [
    expensive_function(i) if expensive_function(i) for i in range(0, 10)
]

Of course, you could use normal for loop syntax but it’s a fair bit more verbose, and for illustrative purposes, we’re allergic to verbose.

y = []
for a in range(0, 10):
    x = expensive_function(i)
    if x == 0:
        y.append(x)

The walrus operator plops to the rescue here and allows us to use a list comprehension.

y = [
    x if (x := expensive_function(i)) for i in range(0, 10)
]

This also applies to dictionary comprehensions.

y = {
    i: x if (x := expensive_function(i)) for i in range(0, 10)
}

Do While Loops

A do-while loop was proposed for Python in PEP 315 but was rejected for not providing a material improvement over the following:

while True:
    x = f(a, b) # setup code
    if not x:   # condition
        break
    # loop body using x

A shortened version of do-while loop can be accomplished by having setup code execute once before the loop and moving the condition into a while loop. However, this is error-prone; x = f(a, b) is duplicated for both the setup code and the loop body, and if it needs changing there are now multiple places that must be updated.

x = f(a, b) # setup code
while x:    # condition
    x = f(a, b)
    # loop body using x

With the walrus operator, it can all be inlined to the while condition.

while x := f(a, b): # setup code and condition
    # loop body using x

Pattern Matching

The walrus operator can also be useful in Pattern Matching. Structural Pattern Matching was introduced in Python 3.10 with PEP 622. If you’re not yet familiar see PEP 363 for a tutorial. The walrus can be useful to inline a function call and store the return value in a variable for use in the cases.

match x := f(a, b):
    case 0:
        # do stuff with x
    case 1:
        # do more stuff with x
    case 2:
        # even more doing with x

An Over-the-Top Overuse Example

While the walrus operator is handy for shaving off a few lines of code, inlining too much can make code difficult to reason about. Use it sparingly, especially with other code-golfing operators. For example, with the ternary operator.

height = get_height(name) if (name := get_name(user_id)) else None

I think this can be okay but I also think it’s clearer written long-form,

if name := get_name(user_id):
    height = get_height(name)
else:
    height = None

It could be formatted over multiple lines so it’s just as readible as a normal if / else and to keep the benefits of the ternary usage by only assigning height once, but it’s now a whopping 5 lines.

height = (
    get_height(name)
    if (name := get_name(user_id))
    else None
)

And remember, just because you can doesn’t mean you should write code like below, if you can avoid it. This example is modified from my own code.

params: Dict[str, Dict[str, float]]

sampler_weights = (
    {
        ModeEnum(mode): weight
        for mode, weight in normalize_weights(weights).items()
    }
    if (weights := params.get("sampler", {}).get("weights"))
    else {"x": 0.5, "y": 0.5}
)

Here, I’ve slapped a dictionary comprehension, a ternary operator, and a walrus operator into the same expression. There’s a lot going on, but it’s formatted over multiple lines to help delineate what’s happening. The variable, params, holds the contents of a configuration .toml that I needed to parse some weights from and convert into an dictionary of {enum: weight}.