traverse_tree.py

"""@package docstring
B-ASIC Operation Tree Traversing Module.
"""

from typing import List, Optional
from collections import deque

from b_asic.operation import Operation


class Traverse:
    """Traverse operation tree."""

    def __init__(self, operation: Operation):
        """Construct a TraverseTree."""
        self._initial_operation = operation

    def _breadth_first_search(self, start: Operation) -> List[Operation]:
        """Use breadth first search to traverse the operation tree."""
        visited: List[Operation] = [start]
        queue = deque([start])
        while queue:
            operation = queue.popleft()
            for n_operation in operation.neighbours:
                if n_operation not in visited:
                    visited.append(n_operation)
                    queue.append(n_operation)

        return visited

    def traverse(self, type_: Optional[Operation] = None) -> List[Operation]:
        """Traverse the the operation tree and return operation where type matches.
        If the type is None then return the entire tree.

        Keyword arguments:
        type_-- the operation type to search for (default None)
        """

        operations: List[Operation] = self._breadth_first_search(self._initial_operation)
        if type_ is not None:
            operations = [oper for oper in operations if isinstance(oper, type_)]

        return operations