zeitbach.com

Assignment expressions in comprehensions

TIL: Since PEP 572 (Python 3.8) you can use assignment expressions (or named expressions) to introduce variables in comprehensions.

I have sometimes abused for loops to bind a value to a local variable in a comprehension. Completely contrived example:

non_empty_words = [
    stripped
    for word in ['hello  ', ' ', 'world']
    for stripped in [word.strip()]
    if stripped
]

Instead of this, we can do the following:

non_empty_words = [
    stripped
    for word in ['hello  ', ' ', 'world']
    if (stripped := word.strip())
]

However, there is a catch!

An assignment expression does not introduce a new scope. In most cases the scope in which the target will be bound is self-explanatory: it is the current scope. […]

There is one special case: an assignment expression occurring in a list, set or dict comprehension or in a generator expression […] binds the target in the containing scope, honoring a nonlocal or global declaration for the target in that scope, if one exists. For the purpose of this rule the containing scope of a nested comprehension is the scope that contains the outermost comprehension. – PEP 572, Scope of the target

In other words, named expressions are leaking out of comprehensions. If you copy the second code block above and execute it in a Python REPL, you’re not creating only one, but two variables in the current scope.

> non_empty_words
['hello', 'world']
> stripped
'world'
> word
NameError: name 'word' is not defined

The loop variable word is local to the comprehension, as you’d expect, the variable introduced by the assignment expression is not. I find this rather unintuitive and ugly. So abusing for loops might still be better if you like comprehensions, but want to avoid surprising bugs.