changed docstring and changed to set in BFS

b_asic/abstract_operation.py

@@ -4,7 +4,7 @@ TODO: More info.
 """
 
 from abc import abstractmethod
-from typing import List, Dict, Optional, Any
+from typing import List, Set, Dict, Optional, Any
 from numbers import Number
 
 from b_asic.port import InputPort, OutputPort
@@ -107,14 +107,8 @@ class AbstractOperation(Operation, AbstractGraphComponent):
 
 		return neighbours
 
-	def traverse(self) -> 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)
-		"""
-
+	def traverse(self) -> Set[Operation]:
+		"""Traverse the operation tree and return a generator with start point in the operation."""
 		return breadth_first_search(self)
 
 	# TODO: More stuff.

b_asic/utilities.py

@@ -3,19 +3,19 @@ B-ASIC Operation Module.
 TODO: More info.
 """
 
-from typing import List
+from typing import Set
 from collections import deque
 
 from b_asic.operation import Operation
 
-def breadth_first_search(start: Operation) -> List[Operation]:
+def breadth_first_search(start: Operation) -> Operation:
     """Use breadth first search to traverse the operation tree."""
-    visited: List[Operation] = [start]
+    visited: Set[Operation] = {start}
     queue = deque([start])
     while queue:
         operation = queue.popleft()
         yield operation
         for n_operation in operation.neighbours:
             if n_operation not in visited:
-                visited.append(n_operation)
+                visited.add(n_operation)
                 queue.append(n_operation)