11
\$\begingroup\$

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())
\$\endgroup\$
  • 1
    \$\begingroup\$ Don't change your code in response to answers. \$\endgroup\$ – Barry Dec 20 '15 at 20:12
  • \$\begingroup\$ I did it because your code is not addressing my issue exactly. Could you then edit your answer? \$\endgroup\$ – aaragon Dec 20 '15 at 20:14
7
\$\begingroup\$

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
    ...
\$\endgroup\$
  • \$\begingroup\$ But using this approach I still don't have the for c in custom: syntax, do I? Nor I can't iterate over lower-dimensional dictionaries? \$\endgroup\$ – aaragon Dec 20 '15 at 19:47
  • \$\begingroup\$ @aaragon You have for c in iter_over_custom(custom). Or you can make custom its own class and define __iter__ as what I wrote iter_over_custom to be (__iter__ can yield). \$\endgroup\$ – Barry Dec 20 '15 at 19:50
  • \$\begingroup\$ I think I like this last suggestion but I'm not sure I understand exactly how to do that. I create a class, say Custom, and I override the __iter__ method and put inside what you wrote. That's it? \$\endgroup\$ – aaragon Dec 20 '15 at 19:53
  • 1
    \$\begingroup\$ @aaragon Yep. __iter__ can return a generator. \$\endgroup\$ – Barry Dec 20 '15 at 19:54
  • 2
    \$\begingroup\$ You could also use .values() instead of .items() to avoid the need for _ -- really just up to you. \$\endgroup\$ – BenC Dec 20 '15 at 20:22

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.