iterating over the values of a list of ordered dictionaries

Question

I am facing the following problem: I have a special data structure that is a dictionary whose keys are integers (dimensions). The values are also dictionaries whose keys are strings (geometric types) and whose values are numpy arrays (connectivities). Something like this:

custom = {1: {'a': np.zeros(10), 'b': np.zeros(100)}, 2:{'c': np.zeros(20), 'd': np.zeros(200)}}

I use this data structure quite a lot in the code, and every time I need to iterate over all rows in the numpy arrays of this data structure (all values of the dictionaries), I have to type:

for d, delem in custom.items():
    for k, v in delem.items():
        for row in v:
            print(row)

And this is way too verbose. So to make things easier to other developers, I am trying to encapsulate this behavior in a custom class so that I can type:

for row in custom:
    print(row)

Also, an important requirement is that all rows have to be traversed "in order". And so I came up with a Test class that derives from list. The class is initialized by passing a dictionary (maybe I can improve this later by just adding the elements directly to it). I have overridden the __iter__ and __next__ methods to provide the iteration type I want. By default, iteration is done through the highest dimension present. But the user can do iteration over lower-dimensional dictionaries and for that I overrode the __call__ method. The technique works, I an iterate as many times as I want over this data structure, and I accomplished the simplicity I wanted. But I'm a bit concerned about performance, as I don't think that calling self[self.d][self.etype][self.idx-1] is very efficient. Maybe there are other improvements to be made as well that I don't see.

import sys
import numpy as np
from collections import OrderedDict

custom = {0: {'n1': np.array([[3]])}, 1: {'l2': np.array([[ 0,  4],
       [ 4,  5],
       [40, 41],
       [41, 42],       [57,  3],
       [57,  3]])}, 2: {'t3x': np.array([[188, 401, 400],
       [188, 187, 401],
       [187, 205, 401],
       [324, 306, 417],
       [306, 305, 417],
       [305, 416, 417]]), 'q3': np.array([[188, 401, 400, 0],
       [188, 187, 401, 0],
       [187, 205, 401, 0],
       [323, 324, 417, 0],
       [324, 306, 417, 0],
       [306, 305, 417, 0],
       [305, 416, 417, 0]])}}


class Test(list):

    def __init__(self, *args, **kwargs):

        super(Test, self).__init__(*args, **kwargs)

        for k,v in args[0].items():
            self[k] = OrderedDict(v)

        self.d = -1
        self.setup()

    def setup(self):
        self.iterator = iter(self[self.d].keys())
        self.etype = next(self.iterator)
        self.idx = 0

    def __iter__(self):
        self.setup()
        return self

    def __next__(self):

        try:
            self.idx += 1
            return self[self.d][self.etype][self.idx-1]

        except IndexError:

            self.etype = next(self.iterator)
            self.idx = 0
            return self[self.d][self.etype][self.idx-1]

    def __call__(self, d):
        self.d = d-1
        return self


def main(argv=()):

    tst = Test(custom)
    print(tst)

    # iterate over the container
    for el in tst:
        print(el)

    # iterate again over the container
    for el in tst:
        print(el)

    # iterate over lower dimension
    for el in tst(-1):
        print(el)

    return 0

if __name__ == "__main__":
    sys.exit(main())

Answer 1

yield!

You're way overthinking the problem. Python has yield. This is one of the coolest things in Python (IMHO). When you want to iterate over a container like this:

for d, delem in custom.items():
    for k, v in delem.items():
        for row in v:
            print(row)

You just write a function that looks pretty much exactly like that:

def iter_over_custom(custom):     ## or a better name
    for _, delem in custom.items():
       for _, v in delem.items():
           for row in v:
               yield row

We can even drop the last loop if you have a sufficiently recent version of Python that supports yield from:

def iter_over_custom(custom):     ## or a better name
    for _, delem in custom.items():
       for _, v in delem.items():
           yield from v

Or you don't actually need .items() since you just need the values:

def iter_over_custom(custom):
    for delem in custom.values():
        for v in delem.values():
            yield from v

That's it. No need for a custom class to effectively reimplement the same. The initial loop could them be written as:

for row in iter_over_custom(custom):
    print(row)

---

Note that if you had your own class, you could write the above in __iter__ too:

class Custom(object):
    ...
    def __iter__(self):
        for delem in self.whatever.values():
            for v in delem.values():
                yield from v
    ...