Coverage for smartmdao / core.py: 100%

1import logging 

2from dataclasses import dataclass, field 

3from typing import Callable, List, Literal 

4 

5from .models import Step 

6from .solvers import Solver, DAGSolver 

7from .visualization import visualize_pipeline 

8from .validation import TypeChecker, StandardTypeChecker, validate_structure, validate_external_inputs 

9 

10# Initialize module-level logger 

11logger = logging.getLogger(__name__) 

12 

13@dataclass 

14class Pipeline: 

15 steps: list[Step] = field(default_factory=list) 

16 solver: Solver = field(default_factory=DAGSolver) 

17 # Static structural validation (producer/consumer type compatibility) always 

18 # runs before the first execution of a given pipeline shape - it's free. 

19 # Runtime per-call validation is opt-in since it adds overhead to every 

20 # step invocation, which matters inside IterativeSolver's convergence loop. 

21 runtime_type_checks: bool = False 

22 type_checker: TypeChecker = field(default_factory=StandardTypeChecker) 

23 _structure_validated: bool = field(default=False, init=False, repr=False, compare=False) 

24 

25 def add(self, fn: Callable, outputs: list[str] = None): 

26 """ 

27 Add a step to the pipeline. 

28 :param fn: The function to execute. 

29 :param outputs: Optional list of variable names this function produces. 

30 """ 

31 step = Step(fn, outputs) 

32 self.steps.append(step) 

33 self._structure_validated = False 

34 logger.debug(f"Added step '{step.name}' to pipeline.") 

35 return self 

36 

37 def step(self, fn: Callable = None, *, outputs: List[str] = None): 

38 """ 

39 Decorator to register a step. 

40 """ 

41 if fn is not None and callable(fn): 

42 self.add(fn, outputs=outputs) 

43 return fn 

44 

45 def wrapper(func): 

46 self.add(func, outputs=outputs) 

47 return func 

48 

49 return wrapper 

50 

51 def run(self, **inputs): 

52 """ 

53 Validates types, then delegates execution to the configured Solver. 

54 """ 

55 logger.info(f"Starting pipeline execution with {len(self.steps)} steps and inputs: {list(inputs.keys())}") 

56 try: 

57 if not self._structure_validated: 

58 validate_structure(self.steps, self.type_checker) 

59 self._structure_validated = True 

60 

61 validate_external_inputs(self.steps, inputs, self.type_checker) 

62 

63 if self.runtime_type_checks: 

64 result = self.solver.solve(self.steps, inputs, type_checker=self.type_checker) 

65 else: 

66 result = self.solver.solve(self.steps, inputs) 

67 

68 logger.info("Pipeline execution completed successfully.") 

69 return result 

70 except Exception as e: 

71 logger.error(f"Pipeline execution failed: {e}") 

72 raise 

73 

74 def visualize(self, 

75 inputs: List[str] = None, 

76 output_path: str = None, 

77 orientation: Literal["TB", "LR"] = "TB", 

78 graph_type: Literal["flow", "bipartite"] = "flow", 

79 view: bool = True): 

80 """ 

81 Generates a Graphviz diagram of the pipeline. 

82 """ 

83 input_set = set(inputs or []) 

84 logger.debug(f"Generating visualization ({graph_type}) for pipeline.") 

85 

86 visualize_pipeline( 

87 steps=self.steps, 

88 inputs=input_set, 

89 output_path=output_path, 

90 orientation=orientation, 

91 graph_type=graph_type, 

92 view=view 

93 )