Tree levels in lazy Python using “Long zip with”

In this blog post about breadth-first traversal, Gibbons talks about implementing BFT with “long zip with” function

lzw :: (a -> a -> a) -> [a] -> [a] -> [a]
lzw f x:xs y:ys = f x y : lzw f xs ys
lzw _ [] ys = ys
lzw _ xs [] = xs

Then for instance

data Tree a = EmpT | TN a (Tree a) (Tree a)
levels :: Tree a -> [[a]]
levels ET = []
levels (TN x tl tr) = [x] : lzw (++) (levels tl) (levels tr)

My question: My version of levels in Python is working, but my lazy version using generators is not. It’s adding an extra [] level at the end. Why?

I consider this practice for using generators and laziness in Python, and I think the problem is some imprecision in how I’m using yield, return, itertools.zip_longest, etc.

An example of what’s going wrong:

>>> my_tree = tree_from_int_segment(10)
>>> pprint(my_tree)

{'left': {'left': {'left': {'left': None, 'right': None, 'val': 0},
                   'right': None,
                   'val': 1},
          'right': {'left': {'left': None, 'right': None, 'val': 3},
                    'right': None,
                    'val': 4},
          'val': 2},
 'right': {'left': {'left': {'left': None, 'right': None, 'val': 6},
                    'right': None,
                    'val': 7},
           'right': {'left': None, 'right': None, 'val': 9},
           'val': 8},
 'val': 5}

>>> tree_levels(my_tree)

[[5], [2, 8], [1, 4, 7, 9], [0, 3, 6]]

>>> list(tree_levels_lazy(my_tree))

[[5], [2, 8], [1, 4, 7, 9], [0, 3, 6], []]

Here’s my code. Thanks in advance!

from itertools import islice
from itertools import zip_longest


def conditional_apply(g, x, y):
    if x is None:
        return y
    elif y is None:
        return x
    else:
        return g(x, y)


def long_zip_with(f, xs, ys):
    zs = list(zip_longest(xs, ys))
    return [conditional_apply(f, x, y) for x, y in zs]


def long_zip_with_lazy(f, xs, ys):
    zs = zip_longest(xs, ys)
    for x, y in zs:
        yield conditional_apply(f, x, y)


def tree_levels(t):
    if t is None:
        return []

    a = t["val"]
    tl = t["left"]
    tr = t["right"]
    return [[a]] + long_zip_with(list.__add__, tree_levels(tl), tree_levels(tr))


def tree_levels_lazy(t):
    if t is None:
        yield []
        return

    a = t["val"]
    tl = t["left"]
    tr = t["right"]
    first_level = [a]

    yield first_level
    for level in long_zip_with_lazy(list.__add__, tree_levels_lazy(tl), tree_levels_lazy(tr)):
        yield level

# Rest is just code for building a tree to work with


def split_at(k, xs):
    it = iter(xs)
    l = list(islice(it, k))
    r = list(it)
    return l, r


def split_list_for_tree(xs):
    if not xs:
        return None
    n = len(xs)
    xsl_p, xsr = split_at(n//2 + 1, xs)
    x = xsl_p.pop()
    return x, xsl_p, xsr


def unfold_tree(seed_to_sprout, init_value):

    def unfolder(seed):
        sprout = seed_to_sprout(seed)
        if sprout is None:
            return None
        x, seed_l, seed_r = sprout
        return {"val": x,
                "left": unfolder(seed_l),
                "right": unfolder(seed_r)}

    return unfolder(init_value)


def tree_from_int_segment(n):
    return unfold_tree(split_list_for_tree, list(range(n)))

Answer

Since the generator version is not concatenating sublists, it should not yield an empty list:

def tree_levels_lazy(t):
    if t is None:
        #yield []  <-- remove this
        return
    ...