Coverage for src / invariant / node.py: 98.25%

1"""Node and SubGraphNode classes representing vertices in the DAG.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass 

6from typing import Any 

7 

8from invariant.params import ref 

9 

10 

11def _collect_refs(value: Any) -> list[ref]: 

12 """Recursively collect all ref() markers from a value.""" 

13 refs: list[ref] = [] 

14 if isinstance(value, ref): 

15 refs.append(value) 

16 elif isinstance(value, dict): 

17 for v in value.values(): 

18 refs.extend(_collect_refs(v)) 

19 elif isinstance(value, list): 

20 for item in value: 

21 refs.extend(_collect_refs(item)) 

22 return refs 

23 

24 

25@dataclass(frozen=True) 

26class Node: 

27 """A vertex in the DAG defining what operation to perform. 

28 

29 Attributes: 

30 op_name: The name of the operation to execute (must be registered). 

31 params: Static parameters for this node (dict of parameter name -> value). 

32 May contain ref() and cel() markers, and ${...} string interpolation. 

33 deps: List of node IDs that this node depends on (upstream dependencies). 

34 cache: When True (default), the node's result is cached unless it depends on 

35 an ephemeral upstream node. When False, the op is always executed, the 

36 result is never stored, and cache bypass cascades to downstream nodes. 

37 """ 

38 

39 op_name: str 

40 params: dict[str, Any] 

41 deps: list[str] 

42 cache: bool = True 

43 

44 def __post_init__(self) -> None: 

45 """Validate node configuration.""" 

46 if not self.op_name: 

47 raise ValueError("op_name cannot be empty") 

48 if not isinstance(self.params, dict): 

49 raise ValueError("params must be a dictionary") 

50 if not isinstance(self.deps, list): 

51 raise ValueError("deps must be a list") 

52 

53 # Validate that all ref() markers reference declared dependencies 

54 self._validate_refs() 

55 

56 def _validate_refs(self) -> None: 

57 """Validate that all ref() markers in params reference declared dependencies.""" 

58 deps_set = set(self.deps) 

59 refs = _collect_refs(self.params) 

60 

61 for ref_marker in refs: 

62 if ref_marker.dep not in deps_set: 

63 raise ValueError( 

64 f"ref('{ref_marker.dep}') in params references undeclared dependency. " 

65 f"Declared deps: {self.deps}. " 

66 f"Add '{ref_marker.dep}' to deps list." 

67 ) 

68 

69 

70@dataclass(frozen=True) 

71class SubGraphNode: 

72 """A vertex that expands to an internal DAG at execution time. 

73 

74 Has deps and params like Node, but carries an internal graph and output node ID 

75 instead of an op_name. The executor runs the internal graph with resolved params 

76 as context and returns the designated output node's artifact. 

77 """ 

78 

79 params: dict[str, Any] 

80 deps: list[str] 

81 graph: dict[str, Node | SubGraphNode] 

82 output: str 

83 

84 def __post_init__(self) -> None: 

85 """Validate SubGraphNode configuration.""" 

86 if not isinstance(self.params, dict): 

87 raise ValueError("params must be a dictionary") 

88 if not isinstance(self.deps, list): 

89 raise ValueError("deps must be a list") 

90 if not isinstance(self.graph, dict): 

91 raise ValueError("graph must be a dictionary") 

92 if self.output not in self.graph: 

93 raise ValueError( 

94 f"output '{self.output}' must be a key in graph. " 

95 f"Graph keys: {list(self.graph.keys())}." 

96 ) 

97 self._validate_refs() 

98 

99 def _validate_refs(self) -> None: 

100 """Validate that all ref() markers in params reference declared dependencies.""" 

101 deps_set = set(self.deps) 

102 refs = _collect_refs(self.params) 

103 for ref_marker in refs: 

104 if ref_marker.dep not in deps_set: 

105 raise ValueError( 

106 f"ref('{ref_marker.dep}') in params references undeclared dependency. " 

107 f"Declared deps: {self.deps}. " 

108 f"Add '{ref_marker.dep}' to deps list." 

109 )