"""@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