Walrus Operator


Martin McBride, 2020-03-04
Tags operator walrus operator v3.8
Categories python language intermediate python

Assignment expressions are a new feature added in Python 3.8. The operator is := which looks like the eyes and tusks of a walrus.

Example

Here is a simple code fragment that prints the length of a string s, but only if the string is longer than 5:

if len(s) > 5:
  print(len(s))

That is ok, but it calls len(s) twice, which isn't ideal. Repeated code is best avoided where possible from the point of view of readability, and of course it can be inefficient to make unnecessary function calls.

The standard way to avoid this would be:

x = len(s)
if x > 5:
  print(x)

This solves the problem but adds an extra line.

The walrus operator allows you to do this:

if (x := len(s)) > 5:
  print(x)

The := operator assigns the value of len(s) to the variable x, but still allows us use the result in the comparison expression.

Assignment expressions

Until 3.8, assignment in Python could only be done via assignment statements like this:

x = len(s)

But this is a statement, so it cannot be used where an expression is required. In other words, you can't do this:

if (x = len(s)) > 5:   # Invalid syntax!
  print(x)

Assignment expressions can be thought of as being expressions that have the side effect of assigning a variable. In other words:

x := len(s)

is exactly like len(s), but it also assigns the result of len(s) to the variable x.

One thing to bear in mind is that assignment operators don't allow you to do anything new in Python. There is nothing you can do with the walrus operator that you couldn't do reasonably easily without it. What it can do in certain circumstances is make your code a little neater and more readable, without needing to resort to repeated calculations.

While loops

Another area where assignment expressions can help if with while loops, for example if you are reading blocks of data from a file object f. Here is how this typically would look:

data = f.read(1024)
while data:
  process(data)
  data = f.read(1024)

The while loop processes data while ever read returns a true value. However, to cope with the fact that the file might be empty, we need an extra read before the loop, and then we need to pre-read the next block of data at the end of each loop. It works, but it is a little clunky.

Here is how we would do it with an assignment expression:

while data := f.read(1024):
  process(data)

This is much neater. At the start of each loop we read the data that we are going to use. We store it in data at the same time as checking if it is false (indicating the end of the file).

List comprehensions

In the next example we will use a list comprehension to calculate sin(x) for each element of an input sequence s, but we wish to filter out any values that are less than 0. We could do it like this:

v = [math.sin(x) for x in s if math.sin(x) >= 0]

This is not ideal, again because we are calling sin(x) twice on the same value. This is both bad style and potentially slow since sin is a relatively expensive function to calculate. Unfortunately, without the walrus operator the only obvious way to fix this was to use two list comprehensions:

temp = [math.sin(x) for x in s]
v = [x for x in temp if x >= 0]

This is also pretty horrible. But here is the solution using a walrus operator:

v = [y for x in s if (y := math.sin(x)) >= 0]

This avoids calling sin(x) twice, and is arguably more declarative than the original case because it is obvious that you are using the same function for both the test and the value transformation from s to v.

Comparison with other languages

Several existing languages already have a similar feature. Perhaps the most well known is C, first implemented in 1972, which has always treated assignments as expressions rather than statements. Of course, there is no shame in pinching a useful feature from one of the classic languages.

So in C, here is a simple assignment statement:

x = 2 + 1;

This statement containing a single expression x = 2 + 1. The expression takes the value 3, the value assigned. This means you can do something like this:

x = y = z = 2 + 1;

To assign the value 3 to all the variables. The calculation 2 + 1 is only performed once.

Assignment also works with if and while constructs, for example:

if (a = f(x))
{
  perform_action(a);
}

This will assign f(x) to a. If the result of calling f(x) is non-zero it will call perform_action(a). Unlike Python, we don't use a special walrus operator, we just use the normal assignment operator =.

This can have negative consequences. Suppose we actually wanted to do this:

if (a == f(x))
{
  perform_action(a);
}

This code calls f(x) and checks if its value is equal to a. It doesn't change a. It will only call perform_action(a) if the value of f(x) is equal to a.

The problem is that the two code blocks shown are both valid, and both useful. Unfortunately they have very different results, but they only differ in one character = instead of ==. This is quite a common source of bugs in C.

In Python, the = operator in a conditional expression would be a syntax error. You can either use := for an assignment, or == for a compare. It is much more difficult to mix the two accidentally.

Limitations

The walrus operator has a fairly limited purpose - it is intended to be created cleaner code when you have a simple expression that is used in a condition (if or while statement) and is also required to be used again soon afterwards.

It doesn't do anything that can't be done in other ways, so the best advice is that if it doesn't make your code clearer, don't use it.


Tag cloud

2d arrays abstract data type alignment and array arrays bezier curve built-in function close closure colour comparison operator comprehension context conversion data types design pattern device space dictionary duck typing efficiency encryption enumerate filter font font style for loop function function composition function plot functools generator gif gradient greyscale higher order function html image processing imagesurface immutable object index inner function input installing iter iterator itertools lambda function len linspace list list comprehension logical operator lru_cache mandelbrot map monad mutability named parameter numeric python numpy object open operator optional parameter or partial application path positional parameter print pure function radial gradient range recipes recursion reduce rgb rotation scaling sequence slice slicing sound spirograph str stream string subpath symmetric encryption template text text metrics transform translation transparency tuple unpacking user space vectorisation webserver website while loop zip

Copyright (c) Axlesoft Ltd 2020