Binary tree implementation with python generators & dataclasses.

binary_tree_traversal.py

from dataclasses import dataclass
from typing import Any


@dataclass
class EmptyNode:
    data: Any

    def __iter__(self):
        return self.__generator__()

    def __generator__(self):
        yield None


@dataclass
class Node:
    data: Any
    left: EmptyNode = EmptyNode()
    right: EmptyNode = EmptyNode()

    def __iter__(self):
        return self.__generator__()

    def __generator__(self):
        yield from iter(self.left)
        yield self.data
        yield from iter(self.right)

    def traverse(self, skip_none=True):
        for element in self:
            if skip_none and element is None:
                continue
            print(element)


class InOrderNode(Node):
    pass


class PostOrderNode(Node):
    def __generator__(self):
        yield from iter(self.left)
        yield from iter(self.right)
        yield self.data


class PreOrderNode(Node):
    def __generator__(self):
        yield self.data
        yield from iter(self.left)
        yield from iter(self.right)


def create_node(traverse_order: str):
    return {
        "preorder": PreOrderNode,
        "postorder": PostOrderNode,
        "inorder": InOrderNode,
    }.get(traverse_order, InOrderNode)


if __name__ == "__main__":
    node = create_node("inorder")
    root = node(4)
    root.left = node(2)
    root.left.left = node(1)
    root.left.right = node(3)
    root.right = node(5)

    root.traverse()