From syntax to success: Mastering the walrus operator for cleaner code

Without assignment expressions, you would have to use two separate statements if you want to bind a value to a name and use that value in an expression. For example, it is quite common to see code similar to:

# walrus.if.py
remainder = value % modulus
if remainder:
  print(f"Not divisible! The remainder is {remainder}."y        

With assignment expressions, we could rewrite this as:

# walrus.if.p
if remainder := value % modulus:
  print(f"Not divisible! The remainder is {remainder}.")y        

Assignment expressions allow us to write fewer lines of code. Used with care, they can also lead to cleaner, more understandable code. Let’s look at a slightly bigger example to see how an assignment expression can really simplify a while loop.

In interactive scripts, we often need to ask the user to choose between a number of options. For example, suppose we are writing an interactive script that allows customers at an ice cream shop to choose what flavor they want. To avoid confusion when preparing orders, we want to ensure that the user chooses one of the available flavors. Without assignment expressions, we might write something like this:

# menu.no.walrus.py
flavors = ["pistachio", "malaga", "vanilla", "chocolate", "strawberry"]
prompt = "Choose your flavor: "
print(flavors)
while True:
  choice = input(prompt)
  if choice in flavors:
  break
  print(f"Sorry, '{choice}' is not a valid option.")
print(f"You chose '{choice}'.")        

Take a moment to read this code carefully. Notice the condition on that loop: while True means “loop forever.” That’s not what we really want, is it? We want to stop looping when the user inputs a valid flavor (choice in flavors). To achieve that, we’ve used an if statement and a break inside the loop. The logic to control the loop is not immediately obvious. In spite of that, this is actually quite a common pattern when the value needed to control the loop can only be obtained inside the loop.

The input() function is very useful in interactive scripts. It allows you to prompt the user for input, and returns it as a string. You can read more about it in the official Python documentation.

How can we improve on this? Let us try to use an assignment expression:

# menu.walrus.py
flavors = ["pistachio", "malaga", "vanilla", "chocolate", "strawberry"]
prompt = "Choose your flavor: "
print(flavors)
while (choice := input(prompt)) not in flavors:
  print(f"Sorry, '{choice}' is not a valid option.")
print(f"You chose '{choice}'.")        

Now the loop’s conditional expression says exactly what we want. That is much easier to understand. The code is also three lines shorter.

Did you notice the parentheses around the assignment expression? We need them because the := operator has lower precedence than the not in operator. Try removing the parentheses and see what happens.

That's it for Today. See you tomorrow.