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

Does there exist a way in Python 2.7+ to make something like the following?

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

I know you can make anything with just ‘if’:

{ something_if_true for key, value in dict_.items() if condition}

Answers:

Thank you for visiting the Q&A section on Magenaut. Please note that all the answers may not help you solve the issue immediately. So please treat them as advisements. If you found the post helpful (or not), leave a comment & I’ll get back to you as soon as possible.

Method 1

You’ve 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.

Method 2

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

Let’s 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 let’s 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'}}

Method 3

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

If you have the same condition for keys and values, you’re 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())

It’s easier to read and the condition is only evaluated once per key, value.

Example with borrowing @Cleb’s 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'}}

Method 4

It also might be worth mentioning that If only statements put the if at the end:

{_ for _ in iterable if True}

Method 5

Another example in using if/else in dictionary comprehension

I am working on data-entry desktop application for my own office work, and it is common for such data-entry application to get all entries from input widget and dump it into a dictionary for further processing like validation, or editing which we must return selected data from file back to entry widgets, etc.

The first round using traditional coding (8 lines):

entries = {'name': 'Material Name', 'maxt': 'Max Working Temperature', 'ther': {100: 1.1, 200: 1.2}}

a_dic, b_dic = {}, {}

for field, value in entries.items():
    if field == 'ther':
        for k,v in value.items():
            b_dic[k] = v
        a_dic[field] = b_dic
    else:
        a_dic[field] = value
    
print(a_dic)
“ {'name': 'Material Name', 'maxt': 'Max Working Temperature', 'ther': {100: 1.1, 200: 1.2}}”

Second round I tried to use dictionary comprehension but the loop still there (6 lines):

entries = {'name': 'Material Name', 'maxt': 'Max Working Temperature', 'ther': {100: 1.1, 200: 1.2}}

for field, value in entries.items():
    if field == 'ther':
        b_dic = {k:v for k,v in value.items()}
        a_dic[field] = b_dic
    else:
        a_dic[field] = value
    
print(a_dic)
“ {'name': 'Material Name', 'maxt': 'Max Working Temperature', 'ther': {100: 1.1, 200: 1.2}}”

Finally, with a one-line dictionary comprehension statement (1 line):

entries = {'name': 'Material Name', 'maxt': 'Max Working Temperature', 'ther': {100: 1.1, 200: 1.2}}

a_dic = {field:{k:v for k,v in value.items()} if field == 'ther' 
        else value for field, value in entries.items()}
    
print(a_dic)
“ {'name': 'Material Name', 'maxt': 'Max Working Temperature', 'ther': {100: 1.1, 200: 1.2}}”

I use python 3.8.3

All methods was sourced from stackoverflow.com or stackexchange.com, is licensed under cc by-sa 2.5, cc by-sa 3.0 and cc by-sa 4.0