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

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

2

3 from __future__ import annotations

4

5 from dataclasses import dataclass

6 from typing import Any

7

8 from invariant.params import ref

9

10

11 def _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)

26 class 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. When False, the op

35 is always executed and the result is never stored (ephemeral node).

36 """

37

38 op_name: str

39 params: dict[str, Any]

40 deps: list[str]

41 cache: bool = True

42

43 def __post_init__(self) -> None:

44 """Validate node configuration."""

45 if not self.op_name:

46 raise ValueError("op_name cannot be empty")

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

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

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

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

51

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

53 self._validate_refs()

54

55 def _validate_refs(self) -> None:

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

57 deps_set = set(self.deps)

58 refs = _collect_refs(self.params)

59

60 for ref_marker in refs:

61 if ref_marker.dep not in deps_set:

62 raise ValueError(

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

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

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

66 )

67

68

69@dataclass(frozen=True)

70 class SubGraphNode:

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

72

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

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

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

76 """

77

78 params: dict[str, Any]

79 deps: list[str]

80 graph: dict[str, Node | SubGraphNode]

81 output: str

82

83 def __post_init__(self) -> None:

84 """Validate SubGraphNode configuration."""

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

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

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

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

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

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

91 if self.output not in self.graph:

92 raise ValueError(

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

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

95 )

96 self._validate_refs()

97

98 def _validate_refs(self) -> None:

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

100 deps_set = set(self.deps)

101 refs = _collect_refs(self.params)

102 for ref_marker in refs:

103 if ref_marker.dep not in deps_set:

104 raise ValueError(

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

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

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

108 )