You Should Be Using Python's Walrus Operator - Here's Why

The assignment operator - or walrus operator as we all know it - is a feature that's been in Python for a while now (since 3.8), yet it's still somewhat controversial and many people have unfounded hate for it.

In this article I will try to convince you that the walrus operator really is a good addition to the language and that if you use it properly, then it can help you make your code more concise and readable.

Basics/Essentials

In case you're not yet familiar with the :=, let's first review some of the basic use cases that might persuade you to give this Python feature a shot.

First example I want to show you is how you can use walrus operator to reduce number of function invocations. Let's imagine a function called func() that performs some very expensive computations. It takes a long time to compute results, so we don't want to call it many times:


# "func" called 3 times
result = [func(x), func(x)**2, func(x)**3]

# Reuse result of "func" without splitting the code into multiple lines
result = [y := func(x), y**2, y**3]

In the first list declaration above, the func(x) is called 3 times, everytime returning same result, which wastes time and compute resources. When rewritten using walrus operator, func() is invoked only once, assigning its result to y and reusing it for remaining 2 list values. You might say, "I can just add y = func(x) before the list declaration and I don't need the walrus!", you can, but that's one extra, unnecessary line of code and at first glance - without knowing that func(x) is super slow - it might not be clear why the extra y variable needs to exist.

If you're not convinced by the above, I have another one. Consider the following list comprehensions with the same expensive func():


result = [func(x) for x in data if func(x)]

result = [y for x in data if (y := func(x))]

In the first line, func(x) is called twice in every loop. Instead - using walrus operator - we compute it once in the if statement and then reuse it. The code length is same, both lines are equally readable, but the second one is twice as efficient. You could avoid using walrus operator while keeping the performance by changing it to full for loop, but that would require 5 lines of code.

One of the most common use cases for walrus operator is reducing nested conditionals, such as when using RegEx matching:


import re

test = "Something to match"

pattern1 = r"^.*(thing).*"
pattern2 = r"^.*(not present).*"

m = re.match(pattern1, test)
if m:
    print(f"Matched the 1st pattern: {m.group(1)}")
else:
    m = re.match(pattern2, test)
    if m:
        print(f"Matched the 2nd pattern: {m.group(1)}")

# ---------------------
# Cleaner
if m := (re.match(pattern1, test)):
    print(f"Matched 1st pattern: '{m.group(1)}'")
elif m := (re.match(pattern2, test)):
    print(f"Matched 2nd pattern: '{m.group(1)}'")

By using walrus, we reduced the matching code from 7 to 4 lines, while making it more readable by removing the nested if.

Next one on the list so-called "loop-and-half" idiom:


while True:  # Loop
    command = input("> ")
    if command == 'exit':  # And a half
        break
    print("Your command was:", command)

# ---------------------
# Cleaner
while (command := input("> ")) != "exit":
    print("Your command was:", command)

The usual solution is to use dummy infinite while loop, with control flow is delegated to the break statement. Instead, we can use walrus operator to reassign the value of command and then use it in while loop's conditional all on the same line, making the code much cleaner and shorter.

Similar simplification can be applied to other while loops as well, for example when reading files line by line or when receiving data from socket.

Accumulate Data In-Place

Moving on to some little more advanced use cases of walrus operator. This one being possibility to accumulate data in-place:


data = [5, 4, 3, 2]
c = 0; print([(c := c + x) for x in data])  # c = 14
# [5, 9, 12, 14]

from itertools import accumulate
print(list(accumulate(data)))

# ---------------------
data = [5, 4, 3, 2]
print(list(accumulate(data, lambda a, b: a*b)))
# [5, 20, 60, 120]

a = 1; print([(a := a*b) for b in data])
# [5, 20, 60, 120]

The first 2 lines show how you can leverage walrus operator to compute running total. For such a simple case, functions from itertools, such as accumulate are a better fit as you can see on the next 2 lines. For more complex scenarios however, using itertools becomes unreadable quite quickly and - in my opinion - version with := is much nicer than the one with lambda.

If you're still not convinced check out the accumulate examples in docs (e.g. the accumulating interest or logistic map example), which aren't really readable. Try rewriting them to use the assignment expression, they will look much nicer.

Naming Values Inside f-string

This example showcases possibilities and limits or the := rather than the best practices.

If you really wanted to, you could use the walrus operator inside f-strings:


from datetime import datetime

print(f"Today is: {(today:=datetime.today()):%Y-%m-%d}, which is {today:%A}")
# Today is: 2022-07-01, which is Friday

from math import radians, sin, cos
angle = 60
print(f'{angle=}\N{degree sign} {(theta := radians(angle)) =: .2f}, {sin(theta) =: .2f}, {cos(theta) =: .2f}')
# angle=60° (theta := radians(angle)) = 1.05, sin(theta) = 0.87, cos(theta) = 0.50

In the first print statement above, we use := to define variable today which is then reused on the same line saving us a repeated call to datetime.today().

Similarly, in the second example, we declare theta variable, which is then reused to compute sin(theta) and cos(theta). In this case we also use it in conjunction with what looks like a "reverse walrus" operator - this really is just = which forces the expression to be printed along its value, plus the : used for formatting the expression.

Notice also that the walrus expressions have to be surrounded by parentheses for f-string to interpret it correctly.

Any and All

You can use Python's any() and all() functions to verify whether any or all values in some iterable satisfy certain condition. What if you however want to also capture the value that caused any() to return True (so-called "witness") or the value that caused all() to fail (so-called "counterexample")?


numbers = [1, 4, 6, 2, 12, 4, 15]

# Only returns boolean, not the values
print(any(number > 10 for number in numbers))  # True
print(all(number < 10 for number in numbers))  # False

# ---------------------
any((value := number) > 10 for number in numbers)  # True
print(value)  # 12

all((counter_example := number) < 10 for number in numbers)  # False
print(counter_example)  # 12

Both any() and all() use short-circuiting to evaluate the expression. This means that they stop the evaluation as soon as they find the first "witness" or "counterexample" respectively. Therefore, with this trick the variable created by walrus operator will always give us the first "witness"/"counterexample".

Gotchas & Limitations

While I tried to motivate you to use walrus operator in previous sections, I think it's also important to warn you about some of its shortcomings and limitations. Following are the gotchas you might run into when using the walrus operator:

In the previous example you saw that short-circuiting can be useful for capturing values in any()/all(), but in some cases it might produce unexpected results:


for i in range(1, 100):
    if (two := i % 2 == 0) and (three := i % 3 == 0):
        print(f"{i} is divisible by 6.")
    elif two:
        print(f"{i} is divisible by 2.")
    elif three:
        print(f"{i} is divisible by 3.")

# NameError: name 'three' is not defined

In the above snippet, we've created a conditional with 2 assignments joined by and which check whether a number is divisible by 2, 3, or 6 based on whether first, second or both conditions are satisfied. At first glance it might seem like a nice trick, but due to short-circuiting, if the expression (two := i % 2 == 0) fails, the second part will be skipped and therefore three will be undefined or will have stale value from previous loop.

Short-circuiting can be beneficial/intended too though. We can use it with regular expressions to search for multiple patterns in a string:


import re

tests = ["Something to match", "Second one is present"]

pattern1 = r"^.*(thing).*"
pattern2 = r"^.*(present).*"

for test in tests:
    m = re.match(pattern1, test)
    if m:
        print(f"Matched the 1st pattern: {m.group(1)}")
    else:
        m = re.match(pattern2, test)
        if m:
            print(f"Matched the 2nd pattern: {m.group(1)}")

# Matched the 1st pattern: thing
# Matched the 2nd pattern: present

for test in tests:
    if m := (re.match(pattern1, test) or re.match(pattern2, test)):
        print(f"Matched: '{m.group(1)}'")
        # Matched: 'thing'
        # Matched: 'present'

We've already seen version of this snippet in the first section where we used if/elif in conjunction with walrus operator. Here we're simplifying it even further by reducing the conditional into single if.

If you're just getting familiar with walrus operator, you might notice that it causes variable scopes to behave differently in comprehensions:


values = [3, 5, 2, 6, 12, 7, 15]

tmp = "unmodified"
dummy = [tmp for tmp in values]
print(tmp)  # As expected, "tmp" was not clobbered - it's still bound to "unmodified"

total = 0
partial_sums = [total := total + v for v in values]
print(total)  # Prints: 50

With normal list/dict/set comprehensions, the loop variable does not leak into the surrounding scope and therefore any existing variables with same name will be unmodified. With walrus operator however, the variable from comprehension (total in the above code) remains accessible after comprehension returns, taking the value from inside comprehension.

When you become more comfortable using walrus in your code, you might try using it in more situations. One place where you should never use it though is with statement:


class ContextManager:
    def __enter__(self):
        print("Entering the context...")

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Leaving the context...")

with ContextManager() as context:
    print(context)  # None

with (context := ContextManager()):
    print(context)  # <__main__.ContextManager object at 0x7fb551cdb9d0>

When using the normal syntax with ContextManager() as context: ..., the context is bound to the return value of context.__enter__(), while if you use the version with :=, then it's bound to the result of ContextManager() itself. This oftentimes doesn't really matter because context.__enter__() usually returns self, but in case it doesn't it will create very hard to debug issues.

For a more practical example see below what happens when you use walrus operator with closing context manager:


from contextlib import closing
from urllib.request import urlopen

with closing(urlopen('https://www.python.org')) as page:
    for line in page:
        print(line)  # Prints website HTML

with (page := closing(urlopen('https://www.python.org'))):
    for line in page:
        print(line)  # TypeError: 'closing' object is not iterable

Another issue you might run into is the relative precedence of :=, which is lower than that of logical operators:


text = "Something to match."
flag = True

if match := re.match(r"^.*(thing).*", text) and flag:
    print(match.groups())  # AttributeError: 'bool' object has no attribute 'group'

if (match := re.match(r"^.*(thing).*", text)) and flag:
    print(match.groups())  # ('thing',)

Here we see that we need to wrap the assignment in parentheses to make sure that result of re.match(...) is assigned to the variable. If we don't, the and expression is evaluated first and boolean result will be assigned instead.

Finally, this really isn't a gotcha, but rather slight limitation. You currently cannot use inline type hints with walrus operator. Therefore, if you want to specify type of the variable, then you need to split it into 2 lines:


from typing import Optional

value: Optional[int] = None
while value := some_func():
    ...  # Do stuff

Closing Thoughts

Like every single other syntax feature, walrus operator can be abused and can decrease clarity and readability. You don't need to shove into you code wherever possible. Treat it as a tool - be aware of its advantages and disadvantages and use it where appropriate.

If you want to see more practical, good usages of walrus operator, check out how it got introduced to the CPython's standard library - all those changes can be found in this PR. Apart from that, I also recommend reading through the PEP 572 which has even more examples as well as rationale for the introduction of the operator.

推荐订阅源

Martin Heinz's Blog