Lazy traversal of a binary tree using a "yield" keyword

btree_traversal.py

from typing import Literal


def inorder_traversal(btree):
    if not btree:
        return

    yield from inorder_traversal(btree[0])
    yield btree[1]
    yield from inorder_traversal(btree[2])


def preorder_traversal(btree):
    if not btree:
        return

    yield btree[1]
    yield from preorder_traversal(btree[0])
    yield from preorder_traversal(btree[2])


def postorder_traversal(btree):
    if not btree:
        return

    yield from postorder_traversal(btree[0])
    yield from postorder_traversal(btree[2])
    yield btree[1]


def btree_walker(btree, order: Literal["inorder", "preorder", "postorder"]):
    """
    Generator that traverses binary tree in the given `order` ('inorder', 'preorder' or 'postorder').
    You should know this from 'Algoritmy II'.
    The tree should be represented with nested tuples (left subtree, value, right subtree).
    If there is no subtree, it will be marked as None.
    Example:
        tree = (((None, 8, None), 3, (None, 4, None)), 5, (None, 1, None))
            5
           / \
          3   1
         / \
        8   4
        list(tree_walker(tree, 'inorder')) == [8, 3, 4, 5, 1]
        list(tree_walker(tree, 'preorder')) == [5, 3, 8, 4, 1]
        list(tree_walker(tree, 'postorder')) == [8, 4, 3, 1, 5]
    """
    match order:
        case "inorder":
            yield from inorder_traversal(btree)
        case "preorder":
            yield from preorder_traversal(btree)
        case "postorder":
            yield from postorder_traversal(btree)
        case _:
            raise ValueError("wrong order")

test_btree_traversal.py

import pytest

from btree_traversal import btree_walker


def test_tree_walker_single():
    assert list(btree_walker((None, 1, None), 'postorder')) == [1]


def test_tree_walker_simple():
    tree = (((None, 8, None), 3, (None, 4, None)), 5, (None, 1, None))
    assert list(btree_walker(tree, 'inorder')) == [8, 3, 4, 5, 1]
    assert list(btree_walker(tree, 'preorder')) == [5, 3, 8, 4, 1]
    assert list(btree_walker(tree, 'postorder')) == [8, 4, 3, 1, 5]


def test_tree_walker_complex():
    tree = (((None, 1, None), 2, ((None, 3, None), 4, (None, 5, None))), 6,
            (None, 7, ((None, 9, None), 8, None)))
    assert list(btree_walker(tree, 'inorder')) == [1, 2, 3, 4, 5, 6, 7, 9, 8]
    assert list(btree_walker(tree, 'preorder')) == [6, 2, 1, 4, 3, 5, 7, 8, 9]
    assert list(btree_walker(tree, 'postorder')) == [1, 3, 5, 4, 2, 9, 8, 7, 6]