Search code examples
python-3.xdictionarygeneratornested-listsyield

Return specific values from a nested dictionary with a generator function


I have a list of dict containing nested dict which looks like below. I would like to use generators and yield to go over reviews list of dict and return all rating given a specific user value. For example "user" : "A11" has two rating records. I would like to have these returned as dict.

data = [{ "product_id" : "123",
    "size" : "L",
    "color" : "blue",
    "reviews" : [
        { "user" : "A11", "rating" : "score1" },
        { "user" : "Z99", "rating" : "score" }] },

    { "product_id" : "987",
    "size" : "M",
    "color" : "red",
    "reviews" : [
        { "user" : "A11", "rating" : "score2" },
        { "user" : "X55", "rating" : "score" }] }
    ]

I have the following generator but it is basically returning all rating values regardless of the user. How can I make the function to filter the values to a specific user ?

def user_rating(nest, kv):
    if isinstance(nest, list):
        for i in nest:
            for x in user_rating(i, kv):
                yield x
    elif isinstance(nest, dict):
        if kv in nest:
            yield nest[kv]
        for j in nest.values():
            for x in user_rating(j, kv):
                yield x

print(list(user_rating(data, 'rating')))

Solution

  • A little modify should be enough:

    def user_rating(nest, kv):
        if isinstance(nest, list):
            for i in nest:
                yield from user_rating(i, kv)
        elif isinstance(nest, dict):
            if nest.get('user') == kv: # <----- here
                yield nest # <---- and here
            for j in nest.values():
                yield from user_rating(j, kv)
                    
    print(list(user_rating(data, 'A11')))
    #[{'user': 'A11', 'rating': 'score1'}, {'user': 'A11', 'rating': 'score2'}]
    

    Btw you can use yield from to ... well, yield from another generator instead of using a for loop