Newer
Older
Jacob Wahlman
committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
"""
B-ASIC Operation Tree Traversing Module.
TODO:
- Get a first operation or? an entire operation tree
- For each start point, follow it to the next operation from it's out port.
- If we are searching for a specific operation end.
- If we are searching for a specific type of operation add the operation to a list and continue.
- When we no more out ports can be traversed return results and end.
"""
from typing import List, Optional
from collections import deque
from b_asic.operation import Operation
class Traverse:
"""Traverse operation tree.
TODO:
- More info.
- Check if a datastructure other than list suits better as return value.
- Implement the type check for operation.
"""
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