python – How can I use if/else in a dictionary comprehension?

python – How can I use if/else in a dictionary comprehension?

Youve already got it: A if test else B is a valid Python expression. The only problem with your dict comprehension as shown is that the place for an expression in a dict comprehension must have two expressions, separated by a colon:

{ (some_key if condition else default_key):(something_if_true if condition
          else something_if_false) for key, value in dict_.items() }

The final if clause acts as a filter, which is different from having the conditional expression.

@Marcins answer covers it all, but just in case someone wants to see an actual example, I add two below:

Lets say you have the following dictionary of sets

d = {key1: {a, b, c}, key2: {foo, bar}, key3: {so, sad}}

and you want to create a new dictionary whose keys indicate whether the string a is contained in the values or not, you can use

dout = {a_in_values_of_{}.format(k) if a in v else a_not_in_values_of_{}.format(k): v for k, v in d.items()}

which yields

{a_in_values_of_key1: {a, b, c},
 a_not_in_values_of_key2: {bar, foo},
 a_not_in_values_of_key3: {sad, so}}

Now lets suppose you have two dictionaries like this

d1 = {bad_key1: {a, b, c}, bad_key2: {foo, bar}, bad_key3: {so, sad}}
d2 = {good_key1: {foo, bar, xyz}, good_key2: {a, b, c}}

and you want to replace the keys in d1 by the keys of d2 if there respective values are identical, you could do

# here we assume that the values in d2 are unique
# Python 2
dout2 = {d2.keys()[d2.values().index(v1)] if v1 in d2.values() else k1: v1 for k1, v1 in d1.items()}

# Python 3
dout2 = {list(d2.keys())[list(d2.values()).index(v1)] if v1 in d2.values() else k1: v1 for k1, v1 in d1.items()}

which gives

{bad_key2: {bar, foo},
 bad_key3: {sad, so},
 good_key2: {a, b, c}}

python – How can I use if/else in a dictionary comprehension?

In case you have different conditions to evaluate for keys and values, @Marcins answer is the way to go.

If you have the same condition for keys and values, youre better off with building (key, value)-tuples in a generator-expression feeding into dict():

dict((modify_k(k), modify_v(v)) if condition else (k, v) for k, v in dct.items())

Its easier to read and the condition is only evaluated once per key, value.

Example with borrowing @Clebs dictionary of sets:

d = {key1: {a, b, c}, key2: {foo, bar}, key3: {so, sad}}

Assume you want to suffix only keys with a in its value and you want the value replaced with the length of the set in such a case. Otherwise, the key-value pair should stay unchanged.

dict((f{k}_a, len(v)) if a in v else (k, v) for k, v in d.items())
# {key1_a: 3, key2: {bar, foo}, key3: {sad, so}}