Coverage for collectra/pipelines/node_graph_manager.py: 92.47%

1"""Node graph management for Collectra workflows. 

2 

3This module provides the NodeGraphManager class for handling node graph operations 

4including node resolution, parent/child data retrieval, and node addition. 

5""" 

6 

7from typing import List, Set, Union 

8 

9import networkx as nx 

10 

11from ..commons.base import NodeStatus 

12from ..tasks.base import Task, TaskNode 

13from ..types.base import Artefact, ArtefactNode, Node 

14 

15 

16class NodeGraphManager: 

17 """Manages node graph operations for the Collectra workflow. 

18 

19 This class encapsulates all graph manipulation operations including: 

20 - Node resolution and retrieval 

21 - Parent and child node queries 

22 - Adding task and artefact nodes to the graph 

23 

24 Attributes: 

25 flow (nx.DiGraph): The NetworkX directed graph representing the workflow. 

26 """ 

27 

28 def __init__(self): 

29 """Initialize the NodeGraphManager with an empty graph.""" 

30 self.flow = nx.DiGraph() 

31 

32 # ========================================================================= 

33 # Graph Lifecycle 

34 # ========================================================================= 

35 

36 def is_empty(self) -> bool: 

37 """Returns True if the graph has no nodes.""" 

38 return len(self.flow.nodes) == 0 

39 

40 def copy_graph(self) -> nx.DiGraph: 

41 """Return a copy of the graph.""" 

42 return self.flow.copy() 

43 

44 # ========================================================================= 

45 # Node Resolution 

46 # ========================================================================= 

47 

48 def resolve_node(self, node_name: str) -> Union[TaskNode, ArtefactNode]: 

49 """Get a node from the workflow graph by name. 

50 

51 Args: 

52 node_name (str): Name of the node to retrieve. 

53 

54 Returns: 

55 Union[TaskNode, ArtefactNode]: The resolved node object. 

56 

57 Raises: 

58 ValueError: If the node is not found or is empty. 

59 """ 

60 node: Union[dict, None] = self.flow.nodes.get(node_name, None) 

61 if not node: 

62 raise ValueError(f"{node_name} not found in workflow") 

63 data: Union[TaskNode, ArtefactNode, None] = node.get("node", None) 

64 if not data: 

65 raise ValueError( 

66 f"data for {node_name} not found in workflow. Possible empty node." 

67 ) 

68 return data 

69 

70 # ========================================================================= 

71 # Node Retrieval (hierarchical: get_nodes -> get_task_nodes / get_artefact_nodes) 

72 # ========================================================================= 

73 

74 def get_nodes(self) -> list[Union[TaskNode, ArtefactNode]]: 

75 """Return all node objects in the graph.""" 

76 return [node["node"] for node in self.flow.nodes.values()] 

77 

78 def get_task_nodes(self) -> list[TaskNode]: 

79 """Return only TaskNode instances from the graph.""" 

80 return [n for n in self.get_nodes() if isinstance(n, TaskNode)] 

81 

82 def get_artefact_nodes(self) -> list[ArtefactNode]: 

83 """Return only ArtefactNode instances from the graph.""" 

84 return [n for n in self.get_nodes() if isinstance(n, ArtefactNode)] 

85 

86 def get_node_names(self) -> list[str]: 

87 """Return all node names in the graph.""" 

88 return list(self.flow.nodes.keys()) 

89 

90 # ========================================================================= 

91 # Parent Retrieval (hierarchical: get_parents -> get_parents_artefact / get_parents_task) 

92 # ========================================================================= 

93 

94 def get_parents(self, node: Node) -> list[Union[TaskNode, ArtefactNode]]: 

95 """Get all resolved parent nodes for a given node.""" 

96 parent_names = list(self.flow.predecessors(str(node.name))) 

97 return [self.resolve_node(name) for name in parent_names] 

98 

99 def get_parents_artefact(self, node: Node) -> List[ArtefactNode]: 

100 """Get all parent artefact nodes for a given node.""" 

101 return [p for p in self.get_parents(node) if isinstance(p, ArtefactNode)] 

102 

103 def get_parents_task(self, node: Node) -> list[TaskNode]: 

104 """Get all parent task nodes for a given node.""" 

105 return [p for p in self.get_parents(node) if isinstance(p, TaskNode)] 

106 

107 def get_ancestor_artefacts(self, node: Node) -> list[ArtefactNode]: 

108 """Return every artefact node upstream of ``node``.""" 

109 return [ 

110 resolved 

111 for name in nx.ancestors(self.flow, str(node.name)) 

112 if isinstance((resolved := self.resolve_node(name)), ArtefactNode) 

113 ] 

114 

115 # ========================================================================= 

116 # Children Retrieval (hierarchical: get_children -> get_children_artefact / get_children_task) 

117 # ========================================================================= 

118 

119 def get_children(self, node: Node) -> list[Union[TaskNode, ArtefactNode]]: 

120 """Get all resolved child nodes for a given node.""" 

121 child_names = list(self.flow.successors(str(node.name))) 

122 return [self.resolve_node(name) for name in child_names] 

123 

124 def get_children_artefact(self, node: Node) -> List[ArtefactNode]: 

125 """Get all child artefact nodes for a given node.""" 

126 return [c for c in self.get_children(node) if isinstance(c, ArtefactNode)] 

127 

128 def get_children_task(self, node: Node) -> list[TaskNode]: 

129 """Get all child task nodes for a given node.""" 

130 return [c for c in self.get_children(node) if isinstance(c, TaskNode)] 

131 

132 # ========================================================================= 

133 # Node Attributes 

134 # ========================================================================= 

135 

136 def set_node_attr(self, node_name: str, **attrs) -> None: 

137 """Set attributes on a graph node.""" 

138 self.flow.nodes[node_name].update(attrs) 

139 

140 # ========================================================================= 

141 # Graph Construction 

142 # ========================================================================= 

143 

144 def has_node(self, node_name: str) -> bool: 

145 """Check if a node exists in the graph.""" 

146 return self.flow.has_node(node_name) 

147 

148 def add_edge(self, source: str, target: str) -> None: 

149 """Add an edge between two nodes.""" 

150 self.flow.add_edge(source, target) 

151 

152 def try_resolve_existing_task(self, name: str, cls_: type) -> Union[Task, None]: 

153 """If node exists and is a TaskNode subclass, return its Task; else None.""" 

154 node_data = self.flow.nodes.get(name, None) 

155 if node_data and issubclass(cls_, Task): 

156 resolved = node_data.get("node", None) 

157 if resolved and isinstance(resolved, TaskNode): 

158 return resolved.get_task() 

159 return None 

160 

161 # ========================================================================= 

162 # Reset 

163 # ========================================================================= 

164 

165 def reset_all_nodes(self) -> None: 

166 """Reset all nodes to their initial state.""" 

167 for node in self.get_nodes(): 

168 if isinstance(node, ArtefactNode): 

169 node.items = dict() 

170 node.ensemble_items = dict() 

171 node.status = NodeStatus.NOT_READY 

172 if isinstance(node, TaskNode): 

173 node.status = NodeStatus.NOT_READY 

174 

175 # ========================================================================= 

176 # Add Nodes 

177 # ========================================================================= 

178 

179 def add_task_node(self, name: str, obj: Task) -> None: 

180 """Add a task node to the workflow graph. 

181 

182 If the node already exists, updates its color attributes. 

183 Otherwise, creates a new TaskNode and adds it to the graph. 

184 

185 Args: 

186 name (str): Name of the task node. 

187 obj (Task): The task object to wrap in the node. 

188 """ 

189 node = self.flow.nodes.get(name, None) 

190 if node: 

191 node["color"] = "blue" 

192 node["fontcolor"] = "white" 

193 return 

194 node = TaskNode(name, obj) 

195 self.flow.add_node( 

196 name, 

197 node=node, 

198 label=str(obj), 

199 shape="box", 

200 color="blue", 

201 fontcolor="white", 

202 style="filled", 

203 ) 

204 

205 def add_artefact_node( 

206 self, 

207 name: str, 

208 obj: Union[Artefact, None] = None, 

209 types: Union[List[type], Set[type]] = [], 

210 ) -> None: 

211 """Add an artefact node to the workflow graph. 

212 

213 If the node already exists, updates it with the new object and types. 

214 Otherwise, creates a new ArtefactNode and adds it to the graph. 

215 

216 Args: 

217 name (str): Name of the artefact node. 

218 obj (Union[Artefact, None], optional): The artefact object to store. Defaults to None. 

219 types (Union[List[type], Set[type]], optional): Expected data types. Defaults to []. 

220 """ 

221 node = self.flow.nodes.get(name, None) 

222 if node: 

223 node["color"] = "salmon" 

224 node["fontcolor"] = "black" 

225 if len(types) == 0 and obj: 

226 types = set([type(obj)]) 

227 if not node: 

228 items = {obj.id: obj} if obj else {} 

229 types = set(types) 

230 node = ArtefactNode(name, items=items, types=types) 

231 self.flow.add_node( 

232 name, 

233 node=node, 

234 label=str(node), 

235 shape="oval", 

236 color="salmon", 

237 fontcolor="black", 

238 style="filled", 

239 ) 

240 elif obj: 

241 artefact_node: ArtefactNode = node["node"] 

242 if not isinstance(artefact_node, ArtefactNode): 

243 return 

244 if type(obj) not in artefact_node.types: 

245 return 

246 artefact_node.add_item(obj) 

247 artefact_node.types = artefact_node.types.union(set(types))