Search code examples
pythondictionarydictionary-comprehension

Check for either key in python dictionary


I have a the following code to get some messages from a telegram bot with python:

useful_messages = [i["message"] for i in response["result"] if i["message"]["from"]["id"] == int(chat_id) and i["message"]["date"] > last_time]

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 1, in <listcomp>
KeyError: 'message'

chat_id is the id of the user I'm interested in, and last_time is used to avoid reading the whole chat with the user.

It worked for some months until today I hit an "edge case":

[i.keys() for i in response["result"]]

[dict_keys(['update_id', 'message']), 
 dict_keys(['update_id', 'message']), 
 dict_keys(['update_id', 'edited_message']), 
 dict_keys(['update_id', 'message'])]

As you can see, one of the messages was edited, and its key is no longer message but edited_message, causing the KeyError above.

I know I can use a for loop, check for either key (message or edited_message) and continue with the message validation (date and id) and extraction. But I wondered if it is possible to check for either key in the dictionary, thus keeping the list/dictionary comprehension (a one-liner solution would be ideal).

I also thought of replacing the edited_message key, if present, by following any of the procedures shown in the answers to this question. Sadly they are hardly one-liners, so the for loop seems to be a less verbose and convoluted solution.

Of course, I'm open to any other solution (or programming logic) that will result in better code.

I'm still new to python, so I'd appreciate your detailed explanations, if complex solutions are offered.


Solution

  • With python>=3.8 you can use the walrus operator

    useful_messages = [x for i in response["result"]
                       if (x := i.get("message", i.get("edited_message")))["from"]["id"] == int(chat_id)
                       and x["date"] > last_time]
    

    Without it would be less performant

    You could do i.get("message", i.get("edited_message")) but keeping the list comprehension would mean repeat it 3 times instead of one, which isn't very performant (unless you don't have many items (not millions))

    useful_messages = [i.get("message", i.get("edited_message")) for i in response["result"] 
                       if i.get("message", i.get("edited_message"))["from"]["id"] == int(chat_id) 
                       and i.get("message", i.get("edited_message"))["date"] > last_time]
    

    A for loop would be cleaner then, a cleaner and performant code values more than a short code

    useful_messages = []
    
    for i in response["result"]:
        msg = i.get("message")
        if msg is None:
            msg = i["edited_message"]
    
        if msg["from"]["id"] == int(chat_id) and msg["date"] > last_time:
            useful_messages.append(msg)