Python’s controversial assignment expression — known as the walrus operator — can improve your code, and it’s time you start using it!
Photo by Jonathan Cooper on UnsplashThe 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.
The first example I want to show you is how you can use the walrus operator to reduce the 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. Here’s how you do this:
# "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 of the declaration above, the func(x) is called three times, every time returning the same result, which wastes time and computing resources. When rewritten using the walrus operator, func() is invoked only once, assigning its result to y and reusing it for the remaining two list values.
You might say, "I can just add y = func(x) before the list declaration, and I don't need the walrus!" Yes, 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 the walrus operator — we compute it once in the if statement and then reuse it. The code length is the same, and both lines are equally readable, but the second one is twice as efficient. You could avoid using the walrus operator while keeping the performance by changing it to a full for loop, but that would require five lines of code.
One of the most common use cases for walrus operator is reducing nested conditionals, such as when using RegEx matching. Here’s the code:
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 seven to four lines while making it more readable by removing the nested if.
The next one on the list is the so-called “loop-and-half” idiom. Here’s what it looks like:
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 a dummy infinite while loop, with control flow delegated to the break statement. Instead, we can use the walrus operator to reassign the value of command and then use it in while loop's conditional on the same line, which will make the code much cleaner and shorter.
A similar simplification can be applied to other while loops as well, for example, when reading files line by line or when receiving data from a socket.
Accumulate Data In-Place
Moving on to some little more advanced use cases of the walrus operator. This one makes it possible to accumulate data in place. Here’s the code:
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 two lines show how you can leverage the walrus operator to compute the running total. For such a simple case, functions from itertools, such as accumulate are a better fit, as shown in the next two lines. For more complex scenarios, however, using itertools becomes unreadable quite quickly. And, in my opinion, the 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 readable. Try rewriting them to use the assignment expression. They will look much nicer.
Naming Values Inside the f-string
This example showcases the possibilities and limits of the := rather than the best practices.
If you wanted to, you could use the walrus operator inside f-strings. Here’s an example:
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.50In 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 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 must be surrounded by parentheses for the 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 conditions. 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) # 12Both any() and all() use short-circuiting to evaluate the expression. This means they stop the evaluation as soon as they find the first "witness" or "counterexample," respectively. Therefore, with this trick, the variable created by the walrus operator will always give us the first "witness"/"counterexample."
Gotchas and Limitations
While I tried to motivate you to use the 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 could be useful for capturing values in any()/ all(), but in some cases, it might produce unexpected results. Here’s an example:
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 definedIn the above snippet, we’ve created a conditional with two assignments joined by and. They check whether a number is divisible by 2, 3, or 6 based on whether the 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 a stale value from the 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, as you can see below:
walrus_short_circuting_regex.py
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 a 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 familiar with the walrus operator, you might notice that it causes variable scopes to behave differently in comprehensions. Here’s the code:
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: 50With normal list/dict/set comprehensions, the loop variable does not leak into the surrounding scope, and therefore, any existing variables with the same name will be unmodified.
With the 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 you should never use it is with a with statement. Here’s the code on this gotcha:
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 often doesn't matter because context.__enter__() usually returns self, but in case it doesn't, it will create hard-to-debug issues.
For a more practical example, see what happens when you use the walrus operator with closing context manager below:
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 iterableAnother issue you might run into is the relative precedence of :=, which is lower than that of logical operators. Here’s the code:
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 a boolean result will be assigned instead.
Finally, this really isn’t a gotcha but rather a slight limitation. You currently cannot use inline-type hints with the walrus operator. Therefore, if you want to specify the type of the variable, then you need to split it into two lines, as shown below:
from typing import Optional
value: Optional[int] = None
while value := some_func():
... # Do stuffClosing Thoughts
Like every other syntax feature, the walrus operator can be abused and decrease clarity and readability. You don’t need to shove it into your 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 the 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 PEP 572, which has even more examples and the rationale for introducing the operator.
Python f-strings Are More Powerful Than You Might Think Learn about the unknown features of Python’s f-strings — the formatted string literals — and up your text formatting…towardsdatascience.com
Automate All the Boring Kubernetes Operations With Python Learn how you can use Python’s Kubernetes Client library to automate all the boring Kubernetes tasks and operationsbetterprogramming.pub
Want to Connect?
This article was originally posted at martinheinz.dev