Coverage for collectra/tasks/base.py: 87.50%

1"""Base task class for Collectra workflows. 

2 

3This module defines the fundamental task interface used throughout the Collectra system. 

4It provides the base classes that all tasks must inherit from and implements the core task execution lifecycle. 

5 

6Classes: 

7 Task: Base class for all workflow tasks 

8""" 

9 

10__all__ = ["Task", "TaskNode"] 

11 

12import re 

13from dataclasses import dataclass 

14from itertools import product 

15from pathlib import Path 

16from typing import TYPE_CHECKING, Generic 

17 

18from collectra.commons import BaseEntity, Node, NodeStatus, T, TaskContext 

19from collectra.cli import command 

20 

21if TYPE_CHECKING: 

22 from collectra.pipelines.base import Collectra 

23 

24 

25class Task(BaseEntity, Generic[T]): 

26 """Abstract base class for all tasks in Collectra workflows. 

27 

28 Defines the common interface and behavior that all tasks must implement. 

29 

30 Attributes: 

31 name (str): Unique identifier for the task within the workflow. 

32 input (list[str]): List of input parameter names. 

33 output (list[str]): List of output names. 

34 context (TaskContext): Runtime context for task execution. 

35 """ 

36 

37 collective: bool = False 

38 

39 def __init__(self, name: str, **kwargs) -> None: 

40 super().__init__(name) 

41 self.pipeline: Collectra | None = None 

42 self.context: TaskContext = TaskContext() 

43 for key, value in kwargs.items(): 

44 setattr(self, key, value) 

45 

46 @property 

47 def input_dict(self) -> dict[str, list]: 

48 input_dict = dict() 

49 if hasattr(self, "input"): 

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

51 input_list = [self.input] 

52 else: 

53 input_list = self.input.copy() 

54 input_dict = {k: [] for k in input_list} 

55 return input_dict 

56 

57 def prepare_inputs(self, parents: list) -> list: 

58 input_dict = self.input_dict 

59 input_keys = list(input_dict.keys()) 

60 for parent in parents: 

61 if parent.name in input_dict: 

62 input_dict[parent.name].extend(parent.items.values()) 

63 

64 if self.collective: 

65 collective_entries = [ 

66 item for values in input_dict.values() for item in values 

67 ] 

68 return [collective_entries] if collective_entries else [] 

69 

70 # Filter out empty inputs (READY but no items) 

71 non_empty_dict = {k: v for k, v in input_dict.items() if v} 

72 if not non_empty_dict: 

73 return [] 

74 non_empty_keys = list(non_empty_dict.keys()) 

75 value_lists = list(non_empty_dict.values()) 

76 entries = [list(combination) for combination in list(product(*value_lists))] 

77 

78 entries_to_remove = [] 

79 for entry_id, entry in enumerate(entries): 

80 entry_items_by_key = {f"{item.name}": item for item in entry} 

81 is_valid = True 

82 for item in entry: 

83 for parent_id in item.parents: 

84 parent_key = re.sub("-.*", "", parent_id) 

85 if ( 

86 parent_key in non_empty_keys 

87 ): # Only validate against non-empty inputs 

88 if ( 

89 parent_key not in entry_items_by_key 

90 or entry_items_by_key[parent_key].id != parent_id 

91 ): 

92 is_valid = False 

93 break 

94 if not is_valid: 

95 entries_to_remove.append(entry_id) 

96 

97 entries = [ 

98 entry for id, entry in enumerate(entries) if id not in entries_to_remove 

99 ] 

100 return entries 

101 

102 def run(self, *args) -> T | None: 

103 """Run the task execution logic. 

104 

105 This method contains the core task execution logic and must be 

106 implemented by all subclasses. 

107 

108 """ 

109 raise NotImplementedError("Subclasses must implement this method.") 

110 

111 def get_output_name(self) -> str: 

112 name = ( 

113 f"{self.get_name()}_output" 

114 if not hasattr(self, "output") 

115 else self.output[0] if isinstance(self.output, list) else self.output 

116 ) 

117 return name 

118 

119 def __call__(self, *args) -> T | None: 

120 """Run the task with the provided arguments. 

121 

122 Args: 

123 **kwargs: Keyword arguments for task execution. 

124 

125 Returns: 

126 

127 """ 

128 return self.run(*args) 

129 

130 @command(name="run") 

131 def cli_run( 

132 self, 

133 inputs: list[str], 

134 output: Path | None = None, 

135 verbose: bool = False, 

136 usage: bool = False, 

137 render: bool = False, 

138 ) -> None: 

139 """Run this task for the supplied input files.""" 

140 if self.pipeline is None: 

141 raise RuntimeError(f"Task {self.name!r} is not attached to a pipeline") 

142 self.pipeline.cli_run( 

143 inputs, 

144 task=self.name, 

145 output=output, 

146 single=True, 

147 verbose=verbose, 

148 usage=usage, 

149 render=render, 

150 ) 

151 

152 

153@dataclass 

154class TaskNode(Node): 

155 

156 task: Task 

157 

158 def get_task(self) -> Task: 

159 return self.task