CODEWITHSUNDEEP
pythonoop
Testuser Semathopy
Testuser Semathopy

Reputation: 85

Iterate over a class composed of other classes

I have a structure representing some data in a following way:

class GenericDataHolder(object):
    def __init__(self):
        self.xxx=[] 
        self.yyy=[]

class A(object):
    def __init__(self):
        self.a1=GenericDataHolder()
        self.a2=GenericDataHolder()

class B(object):
    def __init__(self):
        self.b1=GenericDataHolder()
        self.b2=GenericDataHolder()
        self.b3=A()

I would like to iterate over class B elements so that i will be able to access all nested objects like below:

b=B()
for i in b:
    print i.xxx
    print i.yyy

I have read that there are a few possible solutions to solve this:

This is all fine on basic examples, but I couldn't find any solution to the problem I am facing directly - having class which's elements are composed of other classes.

Upvotes: 3

Views: 105

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1121336

for uses the (internal equivalent of the) iter() function on the object to get an iterator. Implement the __iter__ method to produce one:

class B(object):
    def __init__(self):
        self.b1 = GenericDataHolder()
        self.b2 = GenericDataHolder()
        self.b3 = A()

    def __iter__(self):
        yield self.b1
        yield self.b2
        yield from iter(self.b3)  # use `for sub in self.b3: yield sub` in Python 2

I used a generator function to produce an iterator; when __iter__ is called a generator object is created that then produces b1 and b2 in order.

The (Python 3) yield from syntax then delegates to all such values in b3, which also needs to implement the protocol:

class A(object):
    def __init__(self):
        self.a1 = GenericDataHolder()
        self.a2 = GenericDataHolder()

    def __iter__(self):
        yield self.a1
        yield self.a2

Alternatively, B.__iter__ could just reach into self.b3 and yield self.b3.a1 and self.b3.a2, and not delegate.

Upvotes: 3

Related Questions