Search code examples
pythoniteratorgenerator

Python generator for nested dictionary keys


Let's say my dictionary is

data = {
    'a': {
        'b': {
            'c': {
                'd': 1
            }
        }
    }
}

I want to code something like:

for k,v in iterator(data):
    if some_conditions(k,v):
        v = some_value()

For all the keys of the dictionary, I want to check if the key, value pair matches some condition and update the value in the dictionary if necessary. The idea is that I don't have to keep track of the chain of keys to set the new value.

So, why not an iterator?

The idea seemed good, I would loop through all the keys of the dictionary and update the value on the fly if conditions were met. But I'm having some trouble implementing this, and the reason may be my lack of experience with iterators.

I attempted to code something but it fails and I don't understand why:

def gen(dic):
    for k,v in dic.items():
        if isinstance(v, dict):
            yield v
            gen(v)

g = gen(data)
next(g)

The final data is more complex: some values may be lists having dictionaries as members and I will also have to deal with those dictionaries.

data = {
    'a': {
        'b': {
            'c': [
                {
                    'd': 1,
                    'e': 2
                },
                1,
                {
                    'g': 1,
                    'h': 2
                },
                [
                    {
                        'i': {
                            'j': 1
                        }
                    }
                ]
            ]
        }
    }
}

But if I can figure it out with a simple dictionary, I can then code for the larger case. I am open to different solutions that do not use iterators but remain clean and Pythonic.


Solution

  • You can try:

    class Value:
        def __init__(self, v):
            self.v = v
    
    
    def iterator(o):
        if isinstance(o, dict):
            for k, v in o.items():
                yield k, (val := Value(v))
                o[k] = val.v
                yield from iterator(v)
        if isinstance(o, list):
            for v in o:
                yield from iterator(v)
    
    
    # return True if key is `d` and v is `1`
    def some_condition(k, v):
        return k == "d" and v == 1
    
    
    data = {"a": {"b": {"c": {"d": 1}}}}
    
    for k, val in iterator(data):
        if some_condition(k, val.v):
            val.v = 10
    
    print(data)
    

    Prints:

    {'a': {'b': {'c': {'d': 10}}}}
    

    With the more complicated input:

    data = {"a": {"b": {"c": [{"d": 1, "e": 2}, 1, {"g": 1, "h": 2}, [{"i": {"j": 1}}]]}}}
    
    
    def some_condition(k, v):
        return k == "h" and v == 2
    
    
    for k, val in iterator(data):
        if some_condition(k, val.v):
            val.v = 9999
    
    print(data)
    

    Prints (note the value for the h key is changed):

    {"a": {"b": {"c": [{"d": 1, "e": 2}, 1, {"g": 1, "h": 9999}, [{"i": {"j": 1}}]]}}}