Alternatives to using the walrus operator := inside a nested comprehension

1.4k views Asked by At

Say, for the sake of demonstration, I have to take a list of "input points", and output the three pairs of numbers starting at those points:

1 -> [(1, 2), (3, 4), (5, 6)]
12 -> [(12, 13), (14, 15), (16, 17)]
...

I could do this, in an extended for loop:

points = [1, 12, 42, 69]  # arbitrary inputs
result = []
for point in points:
    rng = range(points, points + 6)
    pairs = list(zip(rng[::2], rng[1::2]))
    result.append(pairs)
return result

I can also do it more concisely with a list comprehension, using the walrus operator to avoid calling range() over and over:

points = [1, 12, 42, 69]
result = [
             list(zip((y := range(p, p + 6))[::2], y[1::2]))
             for p in startpoints
         ]

But if I wanted to, say, use a list comprehension instead of list(), and further process the results of zip() (for example, get my outputs as strings), I can't do that:

points = [1, 12, 42, 69]
result = [
             [
                 str(tup) 
                 for tup in zip((y := range(p, p + 6))[::2], y[1::2])
             ]
             for p in startpoints
         ]
# produces the following error:
  File "<stdin>", line 4
SyntaxError: assignment expression cannot be used in a comprehension iterable expression

PEP 572 mentions problems implementing the assignment operator are at the root of this:

Due to design constraints in the reference implementation (the symbol table analyser cannot easily detect when names are re-used between the leftmost comprehension iterable expression and the rest of the comprehension), named expressions are disallowed entirely as part of comprehension iterable expressions (the part after each "in", and before any subsequent "if" or "for" keyword)

which disappoints me. Are there any workarounds, other than expanding it into a full for loop?

1

There are 1 answers

0
Faboor On

The classic way of doing this is with something like:

[
  [tup for y in [range(p, p + 6)] for tup in zip(y[::2], y[1::2])]
  for p in points
]

However this is rather ugly (and requires constructing additional list/tuple for every p).

You could just avoid the whole range/zip thing with:

[[(y, y + 1) for y in range(p, p + 6, 2)] for p in points]

I believe this makes it more readable (and also slightly faster, because it avoids the zip and 2 out of 3 range calls for each p).