CODEWITHSUNDEEP
pythonfunctionrecursiongenerator
Paradoxis
Paradoxis

Reputation: 4708

Recursively calling an object method that returns an iterator of itself

I'm currently writing a project that requires third party code that uses a method that returns an iterator of itself, an example of how this would look in my code:

def generate():
    for x in obj.children():
        for y in x.children():
            for z in y.children():
                yield z.thing

Currently this simply clutters my code, and becomes hard to read after 3 levels. Ideally I'd get it to do something like this:

x = recursive(obj, method="children", repeat=3).thing

Is there a built in way to do this in Python?

Upvotes: 26

Views: 1742

Answers (3)

Mike Robins
Mike Robins

Reputation: 1773

If using Python 2.7 you need to keep your own stack of iterables and do the looping:

from operator import methodcaller

def recursive(obj, iterater, yielder, depth):
    iterate = methodcaller(iterater)
    xs = [iterate(obj)]
    while xs:
        try:
            x = xs[-1].next()
            if len(xs) != depth:
                xs.append(iterate(x))
            else:
                yield getattr(x, yielder)
        except StopIteration:
            xs.pop()

This a specialized case of a more general recursive ichain from iterable function:

def recursive_ichain(iterable_tree):
    xs = [iter(iterable_tree)]
    while [xs]:
        try:
            x = xs[-1].next()
            if isinstance(x, collections.Iterable):
                xs.append(iter(x))
            else:
                yield x
        except StopIteration:
            xs.pop()

And some test objects:

class Thing(object):
    def __init__(self, thing):
        self.thing = thing

class Parent(object):
    def __init__(self, *kids):
        self.kids = kids

    def children(self):
        return iter(self.kids)

test_obj = Parent(
    Parent(
        Parent(Thing('one'), Thing('two'), Thing('three')),
        Parent(Thing('four')),
        Parent(Thing('five'), Thing('six')),
    ),
    Parent(
        Parent(Thing('seven'), Thing('eight')),
        Parent(),
        Parent(Thing('nine'), Thing('ten')),
    )
)

And testing it:

>>>for t in recursive(test_obj, 'children', 'thing', 3):
>>>    print t
one
two
three
four
five
six
seven
eight
nine
ten

Personnaly I'd be inclined to change the yield getattr(x, yielder) to yield x to access the leaf objects themselves and explicitly access the thing. i.e.

for leaf in recursive(test_obj, 'children', 3):
    print leaf.thing

Upvotes: 7

keredson
keredson

Reputation: 3087

The yield from example above is good, but I seriously doubt the level/depth param is needed. A simpler / more generic solution that works for any tree:

class Node(object):
  def __init__(self, thing, children=None):
    self.thing = thing
    self._children = children
  def children(self):
    return self._children if self._children else []

def generate(node):
  if node.thing:
    yield node.thing
  for child in node.children():
    yield from generate(child)

node = Node('mr.', [Node('derek', [Node('curtis')]), Node('anderson')])
print(list(generate(node)))

Returns:

$ python3 test.py
['mr.', 'derek', 'curtis', 'anderson']

Note this will return the current node's thing before any of its children's. (IE it expresses itself on the way down the walk.) If you'd prefer it to express itself on the way back up the walk, swap the if and the for statements. (DFS vs BFS) But likely doesn't matter in your case (where I suspect a node has either a thing or children, never both).

Upvotes: 5

cs95
cs95

Reputation: 403030

Starting from python3.3, you can use the yield from syntax to yield an entire generator expression.

So, you can modify your function a bit, to take a couple of parameters:

def generate(obj, n):
    if n == 1:
        for x in obj.children():
            yield x.thing
    else:
        for x in obj.children():
            yield from generate(x, n - 1)

The yield from expression will yield the entire generator expression of the recursive call.

Call your function like this:

x = generate(obj, 3)

Note that this returns you a generator of x.things.

Based on your particular requirement, here's a more generic version using getattr that works with arbitrary attributes.

def generate(obj, iterable_attr, attr_to_yield, n):
    if n == 1:
        for x in getattr(obj, iterable_attr):
            yield getattr(x, attr_to_yield)
    else:
        for x in getattr(obj, iterable_attr):
            yield from generate(x, iterable_attr, attr_to_yield, n - 1)

And now, call your function as: 

x = generate(obj, 'children', 'thing', 3)

Upvotes: 28

Related Questions