Coverage for smartmdao / executor.py: 100%

1import inspect 

2import logging 

3from dataclasses import is_dataclass, asdict 

4from typing import Dict, Any, Optional 

5from .models import Step 

6from .validation import TypeChecker, TypeMismatchError 

7 

8# Initialize module-level logger 

9logger = logging.getLogger(__name__) 

10 

11class StepExecutor: 

12 """ 

13 Static helper responsible for binding arguments from memory 

14 and updating memory with results. 

15 """ 

16 @staticmethod 

17 def run_step(step: Step, memory: Dict[str, Any], type_checker: Optional[TypeChecker] = None): 

18 logger.debug(f"Preparing to execute step '{step.name}'") 

19 

20 # Use the robust unwrapped signature to find expected parameters 

21 sig = step.get_signature() 

22 

23 # 1. Bind Arguments 

24 params = {} 

25 missing_required = [] 

26 

27 for name, param in sig.parameters.items(): 

28 if name in memory: 

29 params[name] = memory[name] 

30 elif param.default == inspect.Parameter.empty: 

31 missing_required.append(name) 

32 

33 if missing_required: 

34 error_msg = (f"Step '{step.name}' cannot run. Missing inputs: {missing_required}. " 

35 f"Available in memory: {list(memory.keys())}") 

36 logger.error(error_msg) 

37 raise KeyError(error_msg) 

38 

39 # 1b. Optional Runtime Input Type Check 

40 if type_checker is not None: 

41 StepExecutor._check_input_types(step, params, type_checker) 

42 

43 # 2. Execute 

44 try: 

45 logger.debug(f"Invoking '{step.name}' with inputs: {list(params.keys())}") 

46 result = step.fn(**params) 

47 except Exception as e: 

48 logger.error(f"Error executing step '{step.name}': {e}", exc_info=True) 

49 raise RuntimeError(f"Error executing step '{step.name}': {e}") from e 

50 

51 # 3. Store Result 

52 StepExecutor._update_memory(step, result, memory) 

53 

54 # 3b. Optional Runtime Output Type Check 

55 if type_checker is not None and result is not None: 

56 StepExecutor._check_output_types(step, memory, type_checker) 

57 

58 logger.debug(f"Finished step '{step.name}'.") 

59 

60 @staticmethod 

61 def _check_input_types(step: Step, params: Dict[str, Any], type_checker: TypeChecker): 

62 expected_types = step.resolve_input_types() 

63 for name, value in params.items(): 

64 expected = expected_types.get(name) 

65 if expected is None: 

66 continue 

67 if not type_checker.check_value(value, expected): 

68 raise TypeMismatchError( 

69 f"Step '{step.name}' received {name}={value!r} ({type(value).__name__}), " 

70 f"expected {name}: {getattr(expected, '__name__', expected)}." 

71 ) 

72 

73 @staticmethod 

74 def _check_output_types(step: Step, memory: Dict[str, Any], type_checker: TypeChecker): 

75 # `resolve_output_types()` keys are always a subset of `resolve_output_names()`, 

76 # and `_update_memory` guarantees every one of those names lands in `memory` 

77 # whenever it returns without raising - so a plain lookup is safe here. 

78 expected_types = step.resolve_output_types() 

79 for name, expected in expected_types.items(): 

80 value = memory[name] 

81 if not type_checker.check_value(value, expected): 

82 raise TypeMismatchError( 

83 f"Step '{step.name}' produced {name}={value!r} ({type(value).__name__}), " 

84 f"expected {name}: {getattr(expected, '__name__', expected)}." 

85 ) 

86 

87 @staticmethod 

88 def _update_memory(step: Step, result: Any, memory: Dict[str, Any]): 

89 output_keys = step.resolve_output_names() 

90 

91 if result is None: 

92 logger.debug(f"Step '{step.name}' returned None. No outputs stored.") 

93 return 

94 

95 # Case A: Explicit Manual Outputs (e.g. outputs=['a', 'b']) 

96 if step.manual_outputs: 

97 if len(output_keys) == 1: 

98 memory[output_keys[0]] = result 

99 return 

100 

101 # Handle Dictionary Return with Manual Outputs 

102 if isinstance(result, dict): 

103 for k in output_keys: 

104 if k not in result: 

105 logger.error(f"Step '{step.name}' missing output key '{k}' in returned dict.") 

106 raise KeyError(f"Step '{step.name}' expected output key '{k}' but it was missing in returned dict.") 

107 memory[k] = result[k] 

108 return 

109 

110 # Handle Tuple/List Return with Manual Outputs 

111 if not isinstance(result, (list, tuple)): 

112 raise TypeError(f"Step '{step.name}' expected iterable (or dict) output for keys {output_keys}, got {type(result)}") 

113 

114 if len(result) != len(output_keys): 

115 raise ValueError(f"Step '{step.name}' returned {len(result)} items, expected {len(output_keys)}") 

116 

117 for k, v in zip(output_keys, result): 

118 memory[k] = v 

119 return 

120 

121 # Case B: Dataclass Expansion (Auto-unpacking based on type hint/runtime check) 

122 if is_dataclass(result): 

123 memory.update(asdict(result)) 

124 return 

125 

126 # Case C: Single Default Output 

127 memory[output_keys[0]] = result