How to return an iterable object in Python?

Question

I have a tree. I have several tasks that involving scanning the whole tree.

I would like to write a function, say scanTree () for instance that will scan a tree in a particular order, then returns an iterable object.

The idea is:

def task1 (self):
  x = self.scanTree ()
  for y in x:
    do_something1(y)

def task2 (self):
  x = self.scanTree ()
  for y in x:
    do_something2(y)

But I don't know what kind of object I can return in function scanTree.

Could you give me some hints?

Thank you very much,

Possible duplicate of [Build a Basic Python Iterator](https://stackoverflow.com/questions/19151/build-a-basic-python-iterator) — Mike Scotty, Jun 11 '18 at 09:50
@MikeScotty The accepted answer to that is rather dated, takes an unnecessary detour through the low-level iterator protocol, and misses out on the use of `yield from` which is very useful for recursive traversal. — millimoose, Jun 11 '18 at 09:54

millimoose · Accepted Answer · 2018-06-11T10:04:52.237

This looks like a job for generators. Assuming your btree class looks like this:

class BTree:
  def __init__(self, value, left=None, right=None):
    self._value = value
    self._left = left
    self._right = right

You could then write a scanTree() that does in-order traversal like this:

def scanTree(self):
  '''
  Traverse the binary tree in-order
  '''
  if self._left is not None:
    yield from self._left.scanTree()
  yield self._value
  if self._right is not None:
    yield from self._right.scanTree()

You then use it like this:

>>> l = BTree(1)
>>> r = BTree(3)
>>> t = BTree(2, l, r)
>>> it = t.scanTree()
>>> it
<generator object BTree.scanTree at 0x7fc20602a728>
>>> list(it)
[1, 2, 3]

There's a bunch of concepts involved here to understand this in-depth, but the basic idea is very intuitive: you write the same loop as if you meant to print the elements in the order you want, except you use yield instead of printing, and yield from anytime you recurse.

How to return an iterable object in Python?

1 Answers1