Simplest implementation of tree traversal: PreOrder, InOrder, PostOrder traversal (and basic python decorator)

tree_traversal.py

'''
Implements Tree data structure and few algorithms
'''

def start_end(traverse_func):
    '''
    Decorator
    '''

    def wrapper(self, tree):
        '''
        One can use *args, **kwargs and expand the arguments
        but for now, it's okay to use self and tree
        '''
        print (f'Starting: {traverse_func.__name__}.')
        traverse_func(self, tree)
        print (f'Ending: {traverse_func.__name__}.')
    
    return wrapper

class Node:
    '''
    This code is for class Node
    '''
    def __init__(self, name=None) -> None:
        self.name = name
        self.left = None
        self.right = None
        self.visited = False

    def set_lr(self, left, right):
        '''
        Setter function for binary tree
        '''
        self.left = left
        self.right = right

    @classmethod
    def sample_case1(Cls):
        t_h = Node('H')
        t_i = Node('I')
        t_d = Node('D')
        t_d.set_lr(t_h,t_i)

        t_e = Node('E')
        t_b = Node('B')
        t_b.set_lr(t_d,t_e)

        t_f = Node('F')
        t_g = Node('G')

        t_c = Node('C')
        t_c.set_lr(t_f,t_g)

        t_a = Node('A')
        t_a.set_lr(t_b, t_c)

        return t_a
    
    @classmethod
    def sample_case2(Cls):
        t4 = Node(4)
        t5 = Node(5)
        t2 = Node(2)
        t2.set_lr(t4, t5)
        t3 = Node(3)
        t1 = Node(1)
        t1.set_lr(t2, t3)

        return t1


class Traversal:
    '''
    Traversal class:
    There are different traversal functions defined in this class.
    One should pass a Node object in the traversal functions.
    '''
    def __init__(self) -> None:
        pass

    def pre_order(self, tree):
        # <Root, Left, Right>
        if tree is None:
            return
        print(tree.name)
        # tree.name += " visited"
        tree.visited = True
        self.pre_order(tree.left)
        self.pre_order(tree.right)

    def in_order(self, tree):
        if tree is None:
            return
        self.in_order(tree.left)
        print(tree.name)
        # tree.name += " visited"
        tree.visited = True
        self.in_order(tree.right)

    def post_order(self, tree):
        if tree is None:
            return
        self.post_order(tree.left)
        self.post_order(tree.right)
        print(tree.name)
        # tree.name += " visited"
        tree.visited = True
    
    @start_end
    def bfs(self, tree):
        queue = [tree]
        while len(queue)>0:
            current = queue.pop(0)
            if current is not None:
                print (current.name)
                queue.append(current.left)
                queue.append(current.right)

class TestCode:
    '''
    Driver code for the solution
    '''
    def __init__(self) -> None:
        pass

    @staticmethod
    def run_test():
        traversal = Traversal()
        
        t_a = Node.sample_case1()
        t1 = Node.sample_case2()

        # traversal.in_order(t_a)
        # print (f'----')
        # traversal.post_order(t_a)
        # print (f'----')
        # traversal.pre_order(t_a)
        # print (f'----')
        traversal.bfs(t_a)
        print (f'----')
                
        # traversal.in_order(t1)
        # print (f'----')
        # traversal.pre_order(t1)
        # print (f'----')
        # traversal.post_order(t1)
        # print (f'----')
        traversal.bfs(t1)
        print (f'----')
        

if __name__ == '__main__':
    test_code = TestCode()
    test_code.run_test()