paul23
paul23

Reputation: 9435

Make class iterable respecting inheritance

Well simply making a class iterable is easy enough using meta classes (so some other answers here). However I wish to make a class iterable, and also enabling one to "iterate a subgroup based on inheritance". An example of my use:

class IterPartRegistry(type):
    def __iter__(cls):
        return iter(cls._registry)


class A(object, metaclass=IterPartRegistry):
    _registry = []
    def __init__(self, name):
        self.name = name
        self._registry.append(self)

class B(A):
    pass

class C(A):
    pass


A("A - first")
B("B - first")
B("B - second")
C("C - first")

for t in A:
    print(t.name)

print(" --- ")
for t in B:
    print(t.name)

exit()

The first loop works - it iterates over all instances and childs of "A". However the second loop should only run over the specific subgroup of "A" - those that are instances of the child "B" (or children further down the line).

(How) can this be achieved easiest? In such a way adding more subclasses require least amount of work/change?

Upvotes: 5

Views: 462

Answers (2)

unutbu
unutbu

Reputation: 879601

You could let each class maintain its own list of instances by giving each its own _registry class attribute. Then, instead of checking if each instance is of a particular class, you could instead iterate over all the values in the _registrys for each subclass of cls. To find those subclasses you could use the cls.__subclasses__() method:

import itertools as IT
class IterPartRegistry(type):
    def __init__(cls, name, bases, attrs):
        super(IterPartRegistry, cls).__init__(name, bases, attrs)
        cls._registry = []
    def __iter__(cls):
        yield from cls._registry
        for subcls in cls.__subclasses__():
            yield from subcls

class A(object, metaclass=IterPartRegistry):
    def __init__(self, name):
        self.name = name
        self._registry.append(self)

class B(A): pass

class C(A): pass

class D(B, C): pass

A("A - first")
B("B - first")
B("B - second")
C("C - first")
D("D - first")

for t in A:
    print(t.name)

print(" --- ")
for t in B:
    print(t.name)

yields

A - first
B - first
B - second
D - first
C - first
D - first
 --- 
B - first
B - second
D - first

Upvotes: 2

Yoav Glazner
Yoav Glazner

Reputation: 8066

You can use isinstance to insure that you are getting only class instances

In your code its a one line change:

class IterPartRegistry(type):
    def __iter__(cls):
        return (c for c in cls._registry if isinstance(c, cls))

Upvotes: 3

Related Questions