walrus operator in dict comprehension

Question

I wanted to avoid double evaluation of a mean in a dict comprehension, and I tried using the walrus operator:

>>> dic = {"A": [45,58,75], "B": [55,82,80,92], "C": [78,95,90], "D":[98,75]}
>>> q = {x: (mean := (sum(dic[x]) // len(dic[x]))) for x in dic if mean > 65}

but this gave me the following error:

Traceback (most recent call last):
  File "<pyshell#2>", line 1, in <module>
    q = {x: (mean := (sum(dic[x]) // len(dic[x]))) for x in dic if mean > 65}
  File "<pyshell#2>", line 1, in <dictcomp>
    q = {x: (mean := (sum(dic[x]) // len(dic[x]))) for x in dic if mean > 65}
  NameError: name 'mean' is not defined

This error only happens when I try to use the variable, no problem while defining it:

>>> q = {x: (mean := (sum(dic[x]) // len(dic[x]))) for x in dic if (sum(dic[x]) // len(dic[x])) > 65}
>>> mean
86
>>> q
{'B': 77, 'C': 87, 'D': 86}

Why? Where did I get it wrong?

Elah

Answer 1

Your code is roughly equivalent to

q = {}
for x in dic:
    if mean > 65:
        mean := ...
        q[x] = mean

which means you are using mean before assigning it.

You need to move the definition to the if-clause-section of the dict-comprehension.

>>> dic = {"A": [45,58,75], "B": [55,82,80,92], "C": [78,95,90], "D":[98,75]}
>>> q = {x: mean for x in dic if (mean := (sum(dic[x]) // len(dic[x]))) > 65}
>>> q
{'B': 77, 'C': 87, 'D': 86}

This translates to

q = {}
for x in dic:
    if (mean := ...) > 65:
        q[x] = mean

Idol

Answer 2

You need to place the assignment expression in the conditional of the comprehension, not the value component of the dictionary:

dic = {"A": [45,58,75], "B": [55,82,80,92], "C": [78,95,90], "D":[98,75]}
q = {x:mean for x in dic if (mean := (sum(dic[x]) // len(dic[x]))) > 65}

Output:

{'B': 77, 'C': 87, 'D': 86}

Halonna

Maybe you can add the non-comprehension version of this where it is clear why this is the case (the if part is executed before the x: mean part) EDIT: just saw that @timgeb's answer is already doing that :) – Floccus