Saving intermediate results with the := "walrus"

Sometimes we'll have a complex condition where we want to preserve an expensive intermediate result for later use. Imagine a condition that involves a complex calculation; the cost of computing is high measured in time, or input-output, or memory, or network resource use. Resource use defines the cost of computation.

An example includes doing repetitive searches where the result of the search may be either a useful value or a sentinel value indicating that the target was not found. This is common in the Regular Expression (re) package where the match() method either returns a match object or a None object as a sentinel showing the pattern wasn't found. Once this computation is completed, we may have several uses for the result, and we emphatically do not want to perform the computation again.

This is an example where it can be helpful to assign a name to the value of an expression. We'll look at how to use the "assignment expression" or "walrus" operator. It's called the walrus because the assignment expression operator, :=, looks like the face of a walrus to some people.

Getting ready

Here's a summation where – eventually – each term becomes so small that there's no point in continuing to add it to the overall total:

In effect, this is something like the following summation function:

>>> s = sum((1/(2*n+1))**2 for n in range(0, 20_000))

What's not clear is the question of how many terms are required. In the example, we've summed 20,000. But what if 16,000 are enough to provide an accurate answer?

We don't want to write a summation like this:

>>> b = 0
>>> for n in range(0, 20_000):
...     if (1/(2*n+1))**2 >= 0.000_000_001:
...         b = b + (1/(2*n+1))**2

This example repeats an expensive computation, (1/(2*n+1))**2. That's likely to be a waste of time.

How to do it…

1. Isolate an expensive operation that's part of a conditional test. In this example, the variable term is used to hold the expensive result:

   >>> p = 0
   >>> for n in range(0, 20_000):
   ...     term = (1/(2*n+1))**2
   ...     if term >= 0.000_000_001:
   ...         p = p + term
   

2. Rewrite the assignment statement to use the := assignment operator. This replaces the simple condition of the if statement.
3. Add an else condition to break out of the for statement if no more terms are needed. Here's the results of these two steps:

   >>> q = 0
   >>> for n in range(0, 20_000):
   ...     if (term := (1/(2*n+1))**2) >= 0.000_000_001:
   ...         q = q + term
   ...     else:
   ...         break
   

The assignment expression, :=, lets us do two things in the if statement. We can both compute a value and also check to see that the computed value meets some useful criteria. We can provide the computation and the test criteria adjacent to each other.

How it works…

The assignment expression operator, :=, saves an intermediate result. The operator's result value is the same as the right-hand side operand. This means that the expression a + (b:= c+d) has the same as the expression a+(c+d). The difference between the expression a + (b:= c+d) and the expression a+(c+d) is the side-effect of setting the value of the b variable partway through the evaluation.

An assignment expression can be used in almost any kind of context where expressions are permitted in Python. The most common cases are if statements. Another good idea is inside a while condition.

They're also forbidden in a few places. They cannot be used as the operator in an expression statement. We're specifically prohibited from writing a := 2 as a statement: there's a perfectly good assignment statement for this purpose and an assignment expression, while similar in intent, is potentially confusing.

There's more…

We can do some more optimization of our infinite summation example, shown earlier in this recipe. The use of a for statement and a range() object seems simple. The problem is that we want to end the for statement early when the terms being added are so small that they have no significant change in the final sum.

We can combine the early exit with the term computation:

>>> r = 0
>>> n = 0
>>> while (term := (1/(2*n+1))**2) >= 0.000_000_001:
...     r += term
...     n += 1

We've used a while statement with the assignment expression operator. This will compute a value using (1/(2*n+1))**2, and assign this to term. If the value is significant, we'll add it to the sum, r, and increment the value for the n variable. If the value is too small to be significant, the while statement will end.

Here's another example, showing how to compute running sums of a collection of values. This looks forward to concepts in Chapter 4, Built-In Data Structures Part 1: Lists and Sets. Specifically, this shows a list comprehension built using the assignment expression operator:

>>> data = [11, 13, 17, 19, 23, 29]
>>> total = 0
>>> running_sum = [(total := total + d) for d in data]
>>> total
112
>>> running_sum
[11, 24, 41, 60, 83, 112]

We've started with some data, in the data variable. This might be minutes of exercise each day for most of a week. The value of running_sum is a list object, built by evaluating the expression (total := total + d) for each value, d, in the data variable. Because the assignment expression changes the value of the total variable, the resulting list is the result of each new value being accumulated.

See also

• For details on assignment expression, see PEP 572 where the feature was first described: https://www.python.org/dev/peps/pep-0572/